diff --git a/docs/index.rst b/docs/index.rst index 78e4932e35b3..93afb8d63529 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -165,9 +165,11 @@ Table of Contents source/experimental-features/bleeding-edge source/concepts/warp_environments + source/experimental-features/nerd source/experimental-features/ray source/experimental-features/visuo_tactile_sensor source/experimental-features/rlinf_vla_posttraining + source/experimental-features/supermimic source/features/imitation-learning/index .. toctree:: diff --git a/docs/licenses/dependencies/supermimic-license.txt b/docs/licenses/dependencies/supermimic-license.txt new file mode 100644 index 000000000000..46e96afd1585 --- /dev/null +++ b/docs/licenses/dependencies/supermimic-license.txt @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 ETH Zurich, Chong Zhang + +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/docs/source/_static/css/environment-browser.js b/docs/source/_static/css/environment-browser.js index edb297ec0bf1..ae4baef7a1b5 100644 --- a/docs/source/_static/css/environment-browser.js +++ b/docs/source/_static/css/environment-browser.js @@ -18,6 +18,9 @@ ["Isaac-Cartpole", "rl_games,rsl_rl,skrl,sb3,torchrl", "isaacsim_physx,newton_kamino,newton_mjwarp,ovphysx", "", "", "environment-previews/cartpole.jpg", true, {}, {}, "", null, ["||", "isaacsim_physx||", "newton_kamino||", "newton_mjwarp||", "ovphysx||"]], ["Isaac-Cartpole-Camera-Direct", "rl_games,rsl_rl,skrl", "isaacsim_physx,newton_kamino,newton_mjwarp,ovphysx", "isaacsim_rtx,newton_renderer,ovrtx", "albedo,depth,rgb,semantic_segmentation,simple_shading_constant_diffuse,simple_shading_diffuse_mdl,simple_shading_full_mdl", "tasks/classic/cartpole.jpg", false, {"*": ["rgb"], "rl_games": ["depth"]}, {}, "", null, ["||", "||albedo", "||depth", "||rgb", "||semantic_segmentation", "|isaacsim_rtx|", "|isaacsim_rtx|albedo", "|isaacsim_rtx|depth", "|isaacsim_rtx|rgb", "|isaacsim_rtx|semantic_segmentation", "|isaacsim_rtx|simple_shading_constant_diffuse", "|isaacsim_rtx|simple_shading_diffuse_mdl", "|isaacsim_rtx|simple_shading_full_mdl", "|newton_renderer|", "|newton_renderer|albedo", "|newton_renderer|depth", "|newton_renderer|rgb", "|newton_renderer|semantic_segmentation", "|ovrtx|", "|ovrtx|albedo", "|ovrtx|depth", "|ovrtx|rgb", "|ovrtx|semantic_segmentation", "|ovrtx|simple_shading_constant_diffuse", "|ovrtx|simple_shading_diffuse_mdl", "|ovrtx|simple_shading_full_mdl", "isaacsim_physx||", "isaacsim_physx||albedo", "isaacsim_physx||depth", "isaacsim_physx||rgb", "isaacsim_physx||semantic_segmentation", "isaacsim_physx|isaacsim_rtx|", "isaacsim_physx|isaacsim_rtx|albedo", "isaacsim_physx|isaacsim_rtx|depth", "isaacsim_physx|isaacsim_rtx|rgb", "isaacsim_physx|isaacsim_rtx|semantic_segmentation", "isaacsim_physx|isaacsim_rtx|simple_shading_constant_diffuse", "isaacsim_physx|isaacsim_rtx|simple_shading_diffuse_mdl", "isaacsim_physx|isaacsim_rtx|simple_shading_full_mdl", "isaacsim_physx|newton_renderer|", "isaacsim_physx|newton_renderer|albedo", "isaacsim_physx|newton_renderer|depth", "isaacsim_physx|newton_renderer|rgb", "isaacsim_physx|newton_renderer|semantic_segmentation", "newton_kamino||", "newton_kamino||albedo", "newton_kamino||depth", "newton_kamino||rgb", "newton_kamino||semantic_segmentation", "newton_kamino|isaacsim_rtx|", "newton_kamino|isaacsim_rtx|albedo", "newton_kamino|isaacsim_rtx|depth", "newton_kamino|isaacsim_rtx|rgb", "newton_kamino|isaacsim_rtx|semantic_segmentation", "newton_kamino|isaacsim_rtx|simple_shading_constant_diffuse", "newton_kamino|isaacsim_rtx|simple_shading_diffuse_mdl", "newton_kamino|isaacsim_rtx|simple_shading_full_mdl", "newton_kamino|newton_renderer|", "newton_kamino|newton_renderer|albedo", "newton_kamino|newton_renderer|depth", "newton_kamino|newton_renderer|rgb", "newton_kamino|newton_renderer|semantic_segmentation", "newton_kamino|ovrtx|", "newton_kamino|ovrtx|albedo", "newton_kamino|ovrtx|depth", "newton_kamino|ovrtx|rgb", "newton_kamino|ovrtx|semantic_segmentation", "newton_kamino|ovrtx|simple_shading_constant_diffuse", "newton_kamino|ovrtx|simple_shading_diffuse_mdl", "newton_kamino|ovrtx|simple_shading_full_mdl", "newton_mjwarp||", "newton_mjwarp||albedo", "newton_mjwarp||depth", "newton_mjwarp||rgb", "newton_mjwarp||semantic_segmentation", "newton_mjwarp|isaacsim_rtx|", "newton_mjwarp|isaacsim_rtx|albedo", "newton_mjwarp|isaacsim_rtx|depth", "newton_mjwarp|isaacsim_rtx|rgb", "newton_mjwarp|isaacsim_rtx|semantic_segmentation", "newton_mjwarp|isaacsim_rtx|simple_shading_constant_diffuse", "newton_mjwarp|isaacsim_rtx|simple_shading_diffuse_mdl", "newton_mjwarp|isaacsim_rtx|simple_shading_full_mdl", "newton_mjwarp|newton_renderer|", "newton_mjwarp|newton_renderer|albedo", "newton_mjwarp|newton_renderer|depth", "newton_mjwarp|newton_renderer|rgb", "newton_mjwarp|newton_renderer|semantic_segmentation", "newton_mjwarp|ovrtx|", "newton_mjwarp|ovrtx|albedo", "newton_mjwarp|ovrtx|depth", "newton_mjwarp|ovrtx|rgb", "newton_mjwarp|ovrtx|semantic_segmentation", "newton_mjwarp|ovrtx|simple_shading_constant_diffuse", "newton_mjwarp|ovrtx|simple_shading_diffuse_mdl", "newton_mjwarp|ovrtx|simple_shading_full_mdl", "ovphysx||", "ovphysx||albedo", "ovphysx||depth", "ovphysx||rgb", "ovphysx||semantic_segmentation", "ovphysx|newton_renderer|", "ovphysx|newton_renderer|albedo", "ovphysx|newton_renderer|depth", "ovphysx|newton_renderer|rgb", "ovphysx|newton_renderer|semantic_segmentation", "ovphysx|ovrtx|", "ovphysx|ovrtx|albedo", "ovphysx|ovrtx|depth", "ovphysx|ovrtx|rgb", "ovphysx|ovrtx|semantic_segmentation", "ovphysx|ovrtx|simple_shading_constant_diffuse", "ovphysx|ovrtx|simple_shading_diffuse_mdl", "ovphysx|ovrtx|simple_shading_full_mdl"]], ["Isaac-Cartpole-Camera", "rl_games,rsl_rl", "isaacsim_physx,newton_kamino,newton_mjwarp,ovphysx", "isaacsim_rtx,newton_renderer,ovrtx", "albedo,depth,resnet18,rgb,semantic_segmentation,simple_shading_constant_diffuse,simple_shading_diffuse_mdl,simple_shading_full_mdl,theia_tiny", "tasks/classic/cartpole.jpg", false, {"*": ["rgb"], "rsl_rl": ["resnet18", "theia_tiny"]}, {}, "", null, ["||", "||albedo", "||depth", "||resnet18", "||rgb", "||semantic_segmentation", "||theia_tiny", "|isaacsim_rtx|", "|isaacsim_rtx|albedo", "|isaacsim_rtx|depth", "|isaacsim_rtx|resnet18", "|isaacsim_rtx|rgb", "|isaacsim_rtx|semantic_segmentation", "|isaacsim_rtx|simple_shading_constant_diffuse", "|isaacsim_rtx|simple_shading_diffuse_mdl", "|isaacsim_rtx|simple_shading_full_mdl", "|isaacsim_rtx|theia_tiny", "|newton_renderer|", "|newton_renderer|albedo", "|newton_renderer|depth", "|newton_renderer|resnet18", "|newton_renderer|rgb", "|newton_renderer|semantic_segmentation", "|newton_renderer|theia_tiny", "|ovrtx|", "|ovrtx|albedo", "|ovrtx|depth", "|ovrtx|resnet18", "|ovrtx|rgb", "|ovrtx|semantic_segmentation", "|ovrtx|simple_shading_constant_diffuse", "|ovrtx|simple_shading_diffuse_mdl", "|ovrtx|simple_shading_full_mdl", "|ovrtx|theia_tiny", "isaacsim_physx||", "isaacsim_physx||albedo", "isaacsim_physx||depth", "isaacsim_physx||resnet18", "isaacsim_physx||rgb", "isaacsim_physx||semantic_segmentation", "isaacsim_physx||theia_tiny", "isaacsim_physx|isaacsim_rtx|", "isaacsim_physx|isaacsim_rtx|albedo", "isaacsim_physx|isaacsim_rtx|depth", "isaacsim_physx|isaacsim_rtx|resnet18", "isaacsim_physx|isaacsim_rtx|rgb", "isaacsim_physx|isaacsim_rtx|semantic_segmentation", "isaacsim_physx|isaacsim_rtx|simple_shading_constant_diffuse", "isaacsim_physx|isaacsim_rtx|simple_shading_diffuse_mdl", "isaacsim_physx|isaacsim_rtx|simple_shading_full_mdl", "isaacsim_physx|isaacsim_rtx|theia_tiny", "isaacsim_physx|newton_renderer|", "isaacsim_physx|newton_renderer|albedo", "isaacsim_physx|newton_renderer|depth", "isaacsim_physx|newton_renderer|resnet18", "isaacsim_physx|newton_renderer|rgb", "isaacsim_physx|newton_renderer|semantic_segmentation", "isaacsim_physx|newton_renderer|theia_tiny", "newton_kamino||", "newton_kamino||albedo", "newton_kamino||depth", "newton_kamino||resnet18", "newton_kamino||rgb", "newton_kamino||semantic_segmentation", "newton_kamino||theia_tiny", "newton_kamino|isaacsim_rtx|", "newton_kamino|isaacsim_rtx|albedo", "newton_kamino|isaacsim_rtx|depth", "newton_kamino|isaacsim_rtx|resnet18", "newton_kamino|isaacsim_rtx|rgb", "newton_kamino|isaacsim_rtx|semantic_segmentation", "newton_kamino|isaacsim_rtx|simple_shading_constant_diffuse", "newton_kamino|isaacsim_rtx|simple_shading_diffuse_mdl", "newton_kamino|isaacsim_rtx|simple_shading_full_mdl", "newton_kamino|isaacsim_rtx|theia_tiny", "newton_kamino|newton_renderer|", "newton_kamino|newton_renderer|albedo", "newton_kamino|newton_renderer|depth", "newton_kamino|newton_renderer|resnet18", "newton_kamino|newton_renderer|rgb", "newton_kamino|newton_renderer|semantic_segmentation", "newton_kamino|newton_renderer|theia_tiny", "newton_kamino|ovrtx|", "newton_kamino|ovrtx|albedo", "newton_kamino|ovrtx|depth", "newton_kamino|ovrtx|resnet18", "newton_kamino|ovrtx|rgb", "newton_kamino|ovrtx|semantic_segmentation", "newton_kamino|ovrtx|simple_shading_constant_diffuse", "newton_kamino|ovrtx|simple_shading_diffuse_mdl", "newton_kamino|ovrtx|simple_shading_full_mdl", "newton_kamino|ovrtx|theia_tiny", "newton_mjwarp||", "newton_mjwarp||albedo", "newton_mjwarp||depth", "newton_mjwarp||resnet18", "newton_mjwarp||rgb", "newton_mjwarp||semantic_segmentation", "newton_mjwarp||theia_tiny", "newton_mjwarp|isaacsim_rtx|", "newton_mjwarp|isaacsim_rtx|albedo", "newton_mjwarp|isaacsim_rtx|depth", "newton_mjwarp|isaacsim_rtx|resnet18", "newton_mjwarp|isaacsim_rtx|rgb", "newton_mjwarp|isaacsim_rtx|semantic_segmentation", "newton_mjwarp|isaacsim_rtx|simple_shading_constant_diffuse", "newton_mjwarp|isaacsim_rtx|simple_shading_diffuse_mdl", "newton_mjwarp|isaacsim_rtx|simple_shading_full_mdl", "newton_mjwarp|isaacsim_rtx|theia_tiny", "newton_mjwarp|newton_renderer|", "newton_mjwarp|newton_renderer|albedo", "newton_mjwarp|newton_renderer|depth", "newton_mjwarp|newton_renderer|resnet18", "newton_mjwarp|newton_renderer|rgb", "newton_mjwarp|newton_renderer|semantic_segmentation", "newton_mjwarp|newton_renderer|theia_tiny", "newton_mjwarp|ovrtx|", "newton_mjwarp|ovrtx|albedo", "newton_mjwarp|ovrtx|depth", "newton_mjwarp|ovrtx|resnet18", "newton_mjwarp|ovrtx|rgb", "newton_mjwarp|ovrtx|semantic_segmentation", "newton_mjwarp|ovrtx|simple_shading_constant_diffuse", "newton_mjwarp|ovrtx|simple_shading_diffuse_mdl", "newton_mjwarp|ovrtx|simple_shading_full_mdl", "newton_mjwarp|ovrtx|theia_tiny", "ovphysx||", "ovphysx||albedo", "ovphysx||depth", "ovphysx||resnet18", "ovphysx||rgb", "ovphysx||semantic_segmentation", "ovphysx||theia_tiny", "ovphysx|newton_renderer|", "ovphysx|newton_renderer|albedo", "ovphysx|newton_renderer|depth", "ovphysx|newton_renderer|resnet18", "ovphysx|newton_renderer|rgb", "ovphysx|newton_renderer|semantic_segmentation", "ovphysx|newton_renderer|theia_tiny", "ovphysx|ovrtx|", "ovphysx|ovrtx|albedo", "ovphysx|ovrtx|depth", "ovphysx|ovrtx|resnet18", "ovphysx|ovrtx|rgb", "ovphysx|ovrtx|semantic_segmentation", "ovphysx|ovrtx|simple_shading_constant_diffuse", "ovphysx|ovrtx|simple_shading_diffuse_mdl", "ovphysx|ovrtx|simple_shading_full_mdl", "ovphysx|ovrtx|theia_tiny"]], + ["Isaac-Cartpole-NeRD", "rsl_rl", "", "", "", "environment-previews/cartpole.jpg"], + ["Isaac-Cartpole-NeRD-Explore", "rsl_rl", "", "", "", "environment-previews/cartpole.jpg"], + ["Isaac-Cartpole-NeRD-Learned", "rsl_rl", "", "", "", "environment-previews/cartpole.jpg"], ["Isaac-Fourbar-Pole-Swingup", "rsl_rl", "newton_kamino", "", "", "environment-previews/fourbar_pole.jpg", false, {}, {}, "", null, ["||", "newton_kamino||"]], ["Isaac-Humanoid-Direct", "rl_games,rsl_rl,skrl", "isaacsim_physx,newton_mjwarp,ovphysx", "", "", "environment-previews/humanoid.jpg", true, {}, {}, "", null, ["||", "isaacsim_physx||", "newton_mjwarp||", "ovphysx||"]], ["Isaac-Humanoid", "rl_games,rsl_rl,skrl,sb3", "isaacsim_physx,newton_mjwarp,ovphysx", "", "", "environment-previews/humanoid.jpg", true, {}, {}, "", null, ["||", "isaacsim_physx||", "newton_mjwarp||", "ovphysx||"]], @@ -46,11 +49,17 @@ ["Isaac-Reorient-KukaAllegro-Camera", "rsl_rl", "isaacsim_physx,newton_mjwarp,ovphysx", "isaacsim_rtx,newton_renderer,ovrtx", "albedo128,albedo256,albedo64,cube,depth128,depth256,depth64,duo_camera,raycaster_depth128,raycaster_depth256,raycaster_depth64,rgb128,rgb256,rgb64,semantic_segmentation128,semantic_segmentation256,semantic_segmentation64,shapes,simple_shading_constant_diffuse128,simple_shading_constant_diffuse256,simple_shading_constant_diffuse64,simple_shading_diffuse_mdl128,simple_shading_diffuse_mdl256,simple_shading_diffuse_mdl64,simple_shading_full_mdl128,simple_shading_full_mdl256,simple_shading_full_mdl64,single_camera", "environment-previews/kuka_allegro_reorient.jpg", false, {"*": ["rgb64", "shapes", "single_camera"]}, {}, "", null, ["||", "||albedo128", "||albedo256", "||albedo64", "||cube", "||depth128", "||depth256", "||depth64", "||duo_camera", "||raycaster_depth128", "||raycaster_depth256", "||raycaster_depth64", "||rgb128", "||rgb256", "||rgb64", "||semantic_segmentation128", "||semantic_segmentation256", "||semantic_segmentation64", "||shapes", "||single_camera", "|isaacsim_rtx|", "|isaacsim_rtx|albedo128", "|isaacsim_rtx|albedo256", "|isaacsim_rtx|albedo64", "|isaacsim_rtx|cube", "|isaacsim_rtx|depth128", "|isaacsim_rtx|depth256", "|isaacsim_rtx|depth64", "|isaacsim_rtx|duo_camera", "|isaacsim_rtx|rgb128", "|isaacsim_rtx|rgb256", "|isaacsim_rtx|rgb64", "|isaacsim_rtx|semantic_segmentation128", "|isaacsim_rtx|semantic_segmentation256", "|isaacsim_rtx|semantic_segmentation64", "|isaacsim_rtx|shapes", "|isaacsim_rtx|simple_shading_constant_diffuse128", "|isaacsim_rtx|simple_shading_constant_diffuse256", "|isaacsim_rtx|simple_shading_constant_diffuse64", "|isaacsim_rtx|simple_shading_diffuse_mdl128", "|isaacsim_rtx|simple_shading_diffuse_mdl256", "|isaacsim_rtx|simple_shading_diffuse_mdl64", "|isaacsim_rtx|simple_shading_full_mdl128", "|isaacsim_rtx|simple_shading_full_mdl256", "|isaacsim_rtx|simple_shading_full_mdl64", "|isaacsim_rtx|single_camera", "|newton_renderer|", "|newton_renderer|albedo128", "|newton_renderer|albedo256", "|newton_renderer|albedo64", "|newton_renderer|cube", "|newton_renderer|depth128", "|newton_renderer|depth256", "|newton_renderer|depth64", "|newton_renderer|duo_camera", "|newton_renderer|rgb128", "|newton_renderer|rgb256", "|newton_renderer|rgb64", "|newton_renderer|semantic_segmentation128", "|newton_renderer|semantic_segmentation256", "|newton_renderer|semantic_segmentation64", "|newton_renderer|shapes", "|newton_renderer|single_camera", "|ovrtx|", "|ovrtx|albedo128", "|ovrtx|albedo256", "|ovrtx|albedo64", "|ovrtx|cube", "|ovrtx|depth128", "|ovrtx|depth256", "|ovrtx|depth64", "|ovrtx|duo_camera", "|ovrtx|rgb128", "|ovrtx|rgb256", "|ovrtx|rgb64", "|ovrtx|semantic_segmentation128", "|ovrtx|semantic_segmentation256", "|ovrtx|semantic_segmentation64", "|ovrtx|shapes", "|ovrtx|simple_shading_constant_diffuse128", "|ovrtx|simple_shading_constant_diffuse256", "|ovrtx|simple_shading_constant_diffuse64", "|ovrtx|simple_shading_diffuse_mdl128", "|ovrtx|simple_shading_diffuse_mdl256", "|ovrtx|simple_shading_diffuse_mdl64", "|ovrtx|simple_shading_full_mdl128", "|ovrtx|simple_shading_full_mdl256", "|ovrtx|simple_shading_full_mdl64", "|ovrtx|single_camera", "isaacsim_physx||", "isaacsim_physx||albedo128", "isaacsim_physx||albedo256", "isaacsim_physx||albedo64", "isaacsim_physx||cube", "isaacsim_physx||depth128", "isaacsim_physx||depth256", "isaacsim_physx||depth64", "isaacsim_physx||duo_camera", "isaacsim_physx||raycaster_depth128", "isaacsim_physx||raycaster_depth256", "isaacsim_physx||raycaster_depth64", "isaacsim_physx||rgb128", "isaacsim_physx||rgb256", "isaacsim_physx||rgb64", "isaacsim_physx||semantic_segmentation128", "isaacsim_physx||semantic_segmentation256", "isaacsim_physx||semantic_segmentation64", "isaacsim_physx||shapes", "isaacsim_physx||single_camera", "isaacsim_physx|isaacsim_rtx|", "isaacsim_physx|isaacsim_rtx|albedo128", "isaacsim_physx|isaacsim_rtx|albedo256", "isaacsim_physx|isaacsim_rtx|albedo64", "isaacsim_physx|isaacsim_rtx|cube", "isaacsim_physx|isaacsim_rtx|depth128", "isaacsim_physx|isaacsim_rtx|depth256", "isaacsim_physx|isaacsim_rtx|depth64", "isaacsim_physx|isaacsim_rtx|duo_camera", "isaacsim_physx|isaacsim_rtx|rgb128", "isaacsim_physx|isaacsim_rtx|rgb256", "isaacsim_physx|isaacsim_rtx|rgb64", "isaacsim_physx|isaacsim_rtx|semantic_segmentation128", "isaacsim_physx|isaacsim_rtx|semantic_segmentation256", "isaacsim_physx|isaacsim_rtx|semantic_segmentation64", "isaacsim_physx|isaacsim_rtx|shapes", "isaacsim_physx|isaacsim_rtx|simple_shading_constant_diffuse128", "isaacsim_physx|isaacsim_rtx|simple_shading_constant_diffuse256", "isaacsim_physx|isaacsim_rtx|simple_shading_constant_diffuse64", "isaacsim_physx|isaacsim_rtx|simple_shading_diffuse_mdl128", "isaacsim_physx|isaacsim_rtx|simple_shading_diffuse_mdl256", "isaacsim_physx|isaacsim_rtx|simple_shading_diffuse_mdl64", "isaacsim_physx|isaacsim_rtx|simple_shading_full_mdl128", "isaacsim_physx|isaacsim_rtx|simple_shading_full_mdl256", "isaacsim_physx|isaacsim_rtx|simple_shading_full_mdl64", "isaacsim_physx|isaacsim_rtx|single_camera", "isaacsim_physx|newton_renderer|", "isaacsim_physx|newton_renderer|albedo128", "isaacsim_physx|newton_renderer|albedo256", "isaacsim_physx|newton_renderer|albedo64", "isaacsim_physx|newton_renderer|cube", "isaacsim_physx|newton_renderer|depth128", "isaacsim_physx|newton_renderer|depth256", "isaacsim_physx|newton_renderer|depth64", "isaacsim_physx|newton_renderer|duo_camera", "isaacsim_physx|newton_renderer|rgb128", "isaacsim_physx|newton_renderer|rgb256", "isaacsim_physx|newton_renderer|rgb64", "isaacsim_physx|newton_renderer|semantic_segmentation128", "isaacsim_physx|newton_renderer|semantic_segmentation256", "isaacsim_physx|newton_renderer|semantic_segmentation64", "isaacsim_physx|newton_renderer|shapes", "isaacsim_physx|newton_renderer|single_camera", "newton_mjwarp||", "newton_mjwarp||albedo128", "newton_mjwarp||albedo256", "newton_mjwarp||albedo64", "newton_mjwarp||cube", "newton_mjwarp||depth128", "newton_mjwarp||depth256", "newton_mjwarp||depth64", "newton_mjwarp||duo_camera", "newton_mjwarp||raycaster_depth128", "newton_mjwarp||raycaster_depth256", "newton_mjwarp||raycaster_depth64", "newton_mjwarp||rgb128", "newton_mjwarp||rgb256", "newton_mjwarp||rgb64", "newton_mjwarp||semantic_segmentation128", "newton_mjwarp||semantic_segmentation256", "newton_mjwarp||semantic_segmentation64", "newton_mjwarp||shapes", "newton_mjwarp||single_camera", "newton_mjwarp|isaacsim_rtx|", "newton_mjwarp|isaacsim_rtx|albedo128", "newton_mjwarp|isaacsim_rtx|albedo256", "newton_mjwarp|isaacsim_rtx|albedo64", "newton_mjwarp|isaacsim_rtx|cube", "newton_mjwarp|isaacsim_rtx|depth128", "newton_mjwarp|isaacsim_rtx|depth256", "newton_mjwarp|isaacsim_rtx|depth64", "newton_mjwarp|isaacsim_rtx|duo_camera", "newton_mjwarp|isaacsim_rtx|rgb128", "newton_mjwarp|isaacsim_rtx|rgb256", "newton_mjwarp|isaacsim_rtx|rgb64", "newton_mjwarp|isaacsim_rtx|semantic_segmentation128", "newton_mjwarp|isaacsim_rtx|semantic_segmentation256", "newton_mjwarp|isaacsim_rtx|semantic_segmentation64", "newton_mjwarp|isaacsim_rtx|shapes", "newton_mjwarp|isaacsim_rtx|simple_shading_constant_diffuse128", "newton_mjwarp|isaacsim_rtx|simple_shading_constant_diffuse256", "newton_mjwarp|isaacsim_rtx|simple_shading_constant_diffuse64", "newton_mjwarp|isaacsim_rtx|simple_shading_diffuse_mdl128", "newton_mjwarp|isaacsim_rtx|simple_shading_diffuse_mdl256", "newton_mjwarp|isaacsim_rtx|simple_shading_diffuse_mdl64", "newton_mjwarp|isaacsim_rtx|simple_shading_full_mdl128", "newton_mjwarp|isaacsim_rtx|simple_shading_full_mdl256", "newton_mjwarp|isaacsim_rtx|simple_shading_full_mdl64", "newton_mjwarp|isaacsim_rtx|single_camera", "newton_mjwarp|newton_renderer|", "newton_mjwarp|newton_renderer|albedo128", "newton_mjwarp|newton_renderer|albedo256", "newton_mjwarp|newton_renderer|albedo64", "newton_mjwarp|newton_renderer|cube", "newton_mjwarp|newton_renderer|depth128", "newton_mjwarp|newton_renderer|depth256", "newton_mjwarp|newton_renderer|depth64", "newton_mjwarp|newton_renderer|duo_camera", "newton_mjwarp|newton_renderer|rgb128", "newton_mjwarp|newton_renderer|rgb256", "newton_mjwarp|newton_renderer|rgb64", "newton_mjwarp|newton_renderer|semantic_segmentation128", "newton_mjwarp|newton_renderer|semantic_segmentation256", "newton_mjwarp|newton_renderer|semantic_segmentation64", "newton_mjwarp|newton_renderer|shapes", "newton_mjwarp|newton_renderer|single_camera", "newton_mjwarp|ovrtx|", "newton_mjwarp|ovrtx|albedo128", "newton_mjwarp|ovrtx|albedo256", "newton_mjwarp|ovrtx|albedo64", "newton_mjwarp|ovrtx|cube", "newton_mjwarp|ovrtx|depth128", "newton_mjwarp|ovrtx|depth256", "newton_mjwarp|ovrtx|depth64", "newton_mjwarp|ovrtx|duo_camera", "newton_mjwarp|ovrtx|rgb128", "newton_mjwarp|ovrtx|rgb256", "newton_mjwarp|ovrtx|rgb64", "newton_mjwarp|ovrtx|semantic_segmentation128", "newton_mjwarp|ovrtx|semantic_segmentation256", "newton_mjwarp|ovrtx|semantic_segmentation64", "newton_mjwarp|ovrtx|shapes", "newton_mjwarp|ovrtx|simple_shading_constant_diffuse128", "newton_mjwarp|ovrtx|simple_shading_constant_diffuse256", "newton_mjwarp|ovrtx|simple_shading_constant_diffuse64", "newton_mjwarp|ovrtx|simple_shading_diffuse_mdl128", "newton_mjwarp|ovrtx|simple_shading_diffuse_mdl256", "newton_mjwarp|ovrtx|simple_shading_diffuse_mdl64", "newton_mjwarp|ovrtx|simple_shading_full_mdl128", "newton_mjwarp|ovrtx|simple_shading_full_mdl256", "newton_mjwarp|ovrtx|simple_shading_full_mdl64", "newton_mjwarp|ovrtx|single_camera", "ovphysx||", "ovphysx||albedo128", "ovphysx||albedo256", "ovphysx||albedo64", "ovphysx||cube", "ovphysx||depth128", "ovphysx||depth256", "ovphysx||depth64", "ovphysx||duo_camera", "ovphysx||raycaster_depth128", "ovphysx||raycaster_depth256", "ovphysx||raycaster_depth64", "ovphysx||rgb128", "ovphysx||rgb256", "ovphysx||rgb64", "ovphysx||semantic_segmentation128", "ovphysx||semantic_segmentation256", "ovphysx||semantic_segmentation64", "ovphysx||shapes", "ovphysx||single_camera", "ovphysx|newton_renderer|", "ovphysx|newton_renderer|albedo128", "ovphysx|newton_renderer|albedo256", "ovphysx|newton_renderer|albedo64", "ovphysx|newton_renderer|cube", "ovphysx|newton_renderer|depth128", "ovphysx|newton_renderer|depth256", "ovphysx|newton_renderer|depth64", "ovphysx|newton_renderer|duo_camera", "ovphysx|newton_renderer|rgb128", "ovphysx|newton_renderer|rgb256", "ovphysx|newton_renderer|rgb64", "ovphysx|newton_renderer|semantic_segmentation128", "ovphysx|newton_renderer|semantic_segmentation256", "ovphysx|newton_renderer|semantic_segmentation64", "ovphysx|newton_renderer|shapes", "ovphysx|newton_renderer|single_camera", "ovphysx|ovrtx|", "ovphysx|ovrtx|albedo128", "ovphysx|ovrtx|albedo256", "ovphysx|ovrtx|albedo64", "ovphysx|ovrtx|cube", "ovphysx|ovrtx|depth128", "ovphysx|ovrtx|depth256", "ovphysx|ovrtx|depth64", "ovphysx|ovrtx|duo_camera", "ovphysx|ovrtx|rgb128", "ovphysx|ovrtx|rgb256", "ovphysx|ovrtx|rgb64", "ovphysx|ovrtx|semantic_segmentation128", "ovphysx|ovrtx|semantic_segmentation256", "ovphysx|ovrtx|semantic_segmentation64", "ovphysx|ovrtx|shapes", "ovphysx|ovrtx|simple_shading_constant_diffuse128", "ovphysx|ovrtx|simple_shading_constant_diffuse256", "ovphysx|ovrtx|simple_shading_constant_diffuse64", "ovphysx|ovrtx|simple_shading_diffuse_mdl128", "ovphysx|ovrtx|simple_shading_diffuse_mdl256", "ovphysx|ovrtx|simple_shading_diffuse_mdl64", "ovphysx|ovrtx|simple_shading_full_mdl128", "ovphysx|ovrtx|simple_shading_full_mdl256", "ovphysx|ovrtx|simple_shading_full_mdl64", "ovphysx|ovrtx|single_camera"]], ["Isaac-Shadow-Handover-Direct", "rl_games,rsl_rl,skrl", "isaacsim_physx,newton_mjwarp,ovphysx", "", "", "environment-previews/shadow_hand_over.jpg", false, {}, {"skrl": "MAPPO"}, "", null, ["||", "isaacsim_physx||", "newton_mjwarp||", "ovphysx||"]], ["Isaac-Shadow-Handover", "rsl_rl", "isaacsim_physx,newton_mjwarp,ovphysx", "", "randomized", "environment-previews/shadow_hand_over.jpg", false, {}, {}, "", null, ["||", "||randomized", "isaacsim_physx||", "isaacsim_physx||randomized", "newton_mjwarp||", "newton_mjwarp||randomized", "ovphysx||", "ovphysx||randomized"]], + ["Isaac-Velocity-Flat-Anymal-C-NeRD", "rsl_rl", "", "", ""], + ["Isaac-Velocity-Flat-Anymal-C-NeRD-Explore", "rsl_rl", "", "", ""], + ["Isaac-Velocity-Flat-Anymal-C-NeRD-Learned", "rsl_rl", "", "", ""], ["Isaac-Velocity-Flat-AnymalD", "rsl_rl,skrl", "isaacsim_physx,newton_kamino,newton_mjwarp,ovphysx", "", "", "environment-previews/anymal_d_flat.jpg", true, {}, {}, "", null, ["||", "isaacsim_physx||", "newton_kamino||", "newton_mjwarp||", "ovphysx||"]], ["Isaac-Velocity-Flat-Cassie", "rsl_rl,skrl", "isaacsim_physx,newton_kamino,newton_mjwarp,ovphysx", "", "", "environment-previews/cassie_flat.jpg", true, {}, {}, "", null, ["||", "isaacsim_physx||", "newton_kamino||", "newton_mjwarp||", "ovphysx||"]], ["Isaac-Velocity-Flat-G1", "rsl_rl,skrl", "isaacsim_physx,newton_kamino,newton_mjwarp,ovphysx", "", "", "environment-previews/g1_flat.jpg", true, {}, {}, "", null, ["||", "isaacsim_physx||", "newton_kamino||", "newton_mjwarp||", "ovphysx||"]], ["Isaac-Velocity-Flat-H1", "rsl_rl,skrl", "isaacsim_physx,newton_kamino,newton_mjwarp,ovphysx", "", "", "environment-previews/h1_flat.jpg", true, {}, {}, "", null, ["||", "isaacsim_physx||", "newton_kamino||", "newton_mjwarp||", "ovphysx||"]], ["Isaac-Velocity-Flat-UnitreeGo2", "rsl_rl,skrl", "isaacsim_physx,newton_kamino,newton_mjwarp,ovphysx", "", "", "environment-previews/go2_flat.jpg", true, {}, {}, "", null, ["||", "isaacsim_physx||", "newton_kamino||", "newton_mjwarp||", "ovphysx||"]], + ["Isaac-Velocity-Rough-Anymal-C-NeRD", "rsl_rl", "", "", ""], + ["Isaac-Velocity-Rough-Anymal-C-NeRD-Explore", "rsl_rl", "", "", ""], + ["Isaac-Velocity-Rough-Anymal-C-NeRD-Learned", "rsl_rl", "", "", ""], ["Isaac-Velocity-Rough-AnymalD", "rsl_rl,skrl", "isaacsim_physx,newton_kamino,newton_mjwarp,ovphysx", "", "", "tasks/locomotion/anymal_d_rough.jpg", false, {}, {}, "", null, ["||", "isaacsim_physx||", "newton_kamino||", "newton_mjwarp||", "ovphysx||"]], ["Isaac-Velocity-Rough-Cassie", "rsl_rl,skrl", "isaacsim_physx,newton_kamino,newton_mjwarp,ovphysx", "", "", "tasks/cassie_rough.jpg", false, {}, {}, "", null, ["||", "isaacsim_physx||", "newton_kamino||", "newton_mjwarp||", "ovphysx||"]], ["Isaac-Velocity-Rough-G1", "rsl_rl,skrl", "isaacsim_physx,newton_kamino,newton_mjwarp,ovphysx", "", "", "tasks/locomotion/g1_rough.jpg", false, {}, {}, "", null, ["||", "isaacsim_physx||", "newton_kamino||", "newton_mjwarp||", "ovphysx||"]], @@ -91,6 +100,7 @@ ["IsaacContrib-Lift-Cube-Franka-IK-Abs", "", "", "", "", "environment-previews/franka_lift.jpg"], ["IsaacContrib-Lift-Cube-Franka-IK-Rel", "", "", "", "", "environment-previews/franka_lift.jpg"], ["IsaacContrib-Lift-Cube-OpenArm", "rl_games,rsl_rl", "", "", "", "environment-previews/openarm_uni_lift.jpg"], + ["IsaacContrib-Motion-Tracking-G1", "supermimic", "newton_mjwarp", "", "", "", false, {}, {}, "", null, ["||", "newton_mjwarp||"]], ["IsaacContrib-Multitask-Manipulation", "rsl_rl", "", "", "", "environment-previews/multitask_manipulation.jpg"], ["IsaacContrib-Navigation-3DObstacles-ARL-Robot-1", "rl_games,rsl_rl,skrl", "", "", "", "tasks/drone_arl/arl_robot_1_navigation.jpg"], ["IsaacContrib-Navigation-Flat-AnymalC", "rsl_rl,skrl", "", "", "", "environment-previews/anymal_c_nav.jpg"], @@ -285,6 +295,7 @@ ["skrl", "skrl", "skrl"], ["sb3", "sb3", "Stable-Baselines3"], ["rlinf", "rlinf", "RLinf"], + ["supermimic", "supermimic", "SuperMimic"], ]], ]; diff --git a/docs/source/api/lab/isaaclab.managers.rst b/docs/source/api/lab/isaaclab.managers.rst index d378386b86de..251a56cb6654 100644 --- a/docs/source/api/lab/isaaclab.managers.rst +++ b/docs/source/api/lab/isaaclab.managers.rst @@ -228,9 +228,40 @@ Curriculum Manager :members: :exclude-members: __init__ +.. _recorder-capture: + Recorder Manager ---------------- +Tasks configure recorder terms to capture the state, controls and other fields their consumers need. +Terms run at pre/post step and reset callbacks; reset callbacks receive slices or indices directly. +Consumers own subsequent processing, replay and learning. + +Returning a field path and a batched tensor (or nested dictionary of tensors) retains samples. +For online use without retention, publish the batch on the term and return ``(None, None)``. +Consumers access it through ``env.recorder_manager.get_term(name)`` and copy data before the next capture. +``EXPORT_NONE`` disables file export only; disabling the term removes capture itself. + +Fields may have different sample rates. ``buffer_initial_capacity`` sets the initial samples per world +and field; choose it with payload size and episode length in mind. Storage grows without dropping +samples and is reused across resets. + +**Observation migration:** RL post-step callbacks run after rewards and terminations, before reset +and observation updates. ``env.obs_buf`` still contains the preceding policy input. Read physical +state from the scene or call ``env.observation_manager.compute_group("recording", update_history=False)`` +for current observations. Use a separate noiseless group without history or delay to avoid stale +measurements and additional policy noise draws. History is not advanced, but noise and modifiers still +run. Non-RL environments update their normal observations before post-step capture. + +**Storage migration:** ``get_episode(env_id)`` returns an owned tensor snapshot with a leading sample +dimension; reading copies data and may synchronize the device. Modify live recording through +``add_to_episodes`` and ``set_success_to_episodes``, not the snapshot. ``EpisodeData.add()`` and +``pre_export()`` are removed: collect through the manager, or assign completed tensors to +``episode.data``. HDF5 loading and playback keep the same tensor format. + +The unused ``RecorderTerm.close(file_path)`` hook is removed. Dataset handlers finalize files; +applications explicitly close any additional resources they own. + .. autoclass:: RecorderManager :members: :inherited-members: diff --git a/docs/source/experimental-features/nerd.rst b/docs/source/experimental-features/nerd.rst new file mode 100644 index 000000000000..7ce9fc58a806 --- /dev/null +++ b/docs/source/experimental-features/nerd.rst @@ -0,0 +1,252 @@ +.. _nerd: + +Offline NeRD: Cartpole and ANYmal C +=================================== + +The registered tasks live in ``isaaclab_tasks.core.cartpole.cartpole_nerd_env_cfg`` and +``isaaclab_tasks.core.velocity.config.anymal_c.nerd_env_cfg``. Each has a numerical teacher, +a task-independent ``-Explore`` configuration, and a ``-Learned`` configuration for policy +training/playback. They use ordinary action, observation, reward and termination managers. +The exploration tasks configure ``NeRDRecorderCfg``; teacher and learned playback leave capture +disabled by default. Algorithm settings are registered under ``nerd_cfg_entry_point`` and are +selected from the recording's task metadata, without launching a simulation during offline fitting. + +The shared Newton representation and solver live in ``isaaclab_newton.physics``. The recorder +term lives in ``isaaclab_newton.envs.mdp.recorders``. ``isaaclab_rl.nerd`` owns offline window +sampling and optimization. The two scripts under ``scripts/learning/nerd`` are composition roots +for collection and fitting; policy learning uses the existing reinforcement-learning scripts. +There is no additional integration package or external task launcher. + +PhysicsNeMo Robotics owns ``JointStateCodec`` (including heading conversion), ``GravityEncoder``, +contact encoders, ``FieldNormalizer``, ``LearnedDynamicsSolver`` and its ``DynamicsInferenceSession``. +The native solver adds checkpoint construction and the applied-effort/physical-identity guards. +ParallelNeRD uses its resident temporal cache and repeated-frame selected-world resets; models without +that reset capability use the shared full-window session. Isaac Lab's Newton manager captures the +whole physics step in a Warp CUDA graph, including collision where needed, state writeback and FK. +There is no second Isaac Lab history buffer or inference cache. + +The native HDF5 reader, train/validation split, optimizer schedule and durable W&B publication remain +in the application. This is the native Isaac Lab checkpoint workflow; it does not claim to export +PhysicsNeMo's separate portable ``DynamicsArtifact`` format. + +NeRD physics always consumes applied joint efforts from ``Control.joint_f``. The actuator owns +position/velocity commands, delays, internal state and motor-effort clipping. It runs before each +physics step, and capture copies its final effort at that same solver boundary. There is no +target-conditioned physics mode or recorder control selector. Solver-side motor drives, including +damping drives, and MuJoCo-specific actuators are rejected because they bypass this boundary. +External body wrenches remain unsupported; passive physical forces are part of the fixed model. + +Cartpole records at 120 Hz, with its damping and effort limits implemented by a Newton actuator. +ANYmal C records every 5 ms after ANYdrive evaluates and clips its effort. Its state features use +the root heading frame, body-frame gravity and four fixed foot-sphere probes. The learned task +keeps ANYdrive and advances learned physics four times per 20 ms policy action. Each NeRD step +predicts one physics transition, then Newton forward kinematics updates body state. + +Capture requires ``decimation=1`` and ``num_substeps=1``; policy training can use decimation to +repeat actuator evaluation and learned physics. Both numerical and learned physics enable CUDA +graphs. The learned manager retains a private Torch allocation pool for its captured temporaries. +The existing ANYdrive implementation steps outside the graph, with learned physics replayed at +each 5 ms tick. Position and velocity controllers are supported when their efforts are computed +upstream; NeRD does not predict motor effort as an additional output. + +Rough terrain and live contacts +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +``Isaac-Velocity-Rough-Anymal-C-NeRD`` and its ``-Explore`` / ``-Learned`` variants use the +same 5 ms applied-effort contract with Newton collision detection enabled. NeRD replaces the +numerical dynamics solve; collision detection, ANYdrive and forward kinematics remain. +The exploration task uses terrain-aware reset patches to avoid spawning inside raised terrain. +The policy observations are proprioceptive; these tasks do not add a terrain-height observation. + +Set ``NeRDRecorderCfg.max_contacts`` to a positive capacity to select PhysicsNeMo Robotics' +``RigidContactSetEncoder``. It captures the collision pipeline's pre-solver contacts after +actuator evaluation, without rerunning collision detection. Each packed contact contains validity, +body/other-body identity, dynamic-body status, point, normal, lever arm, signed gap and relative +velocity. The learner converts these to typed, permutation-invariant contact sets: only valid +continuous geometry is normalized; identities and masks retain their semantics. Contact and state +histories use the same episode boundaries and repeated reset prefixes. Capacity overflow stops +capture or deployment instead of silently dropping constraints. + +Terrain geometry can change between collection and deployment because it is represented through +live contacts. The checkpoint still fixes robot geometry, inertias, gains, limits, collision filtering, +gravity and the uniform external contact material. This contract does not support independently +moving external objects or per-surface material variation. Wrench accuracy is a separate validation +requirement; enabling contact conditioning does not enable learned force sensors. + +.. code-block:: bash + + uv run --no-sync python scripts/learning/nerd/record.py \ + --task Isaac-Velocity-Rough-Anymal-C-NeRD-Explore \ + --output logs/nerd/rough/exploration.hdf5 --num_envs 128 --steps 1600 \ + --action_scale 0.75 --action_hold 4 + uv run --no-sync python scripts/learning/nerd/train.py \ + --dataset logs/nerd/rough/exploration.hdf5 --output logs/nerd/rough/model \ + --wandb_project --wandb_entity + +``--action_hold`` sets the interval between action updates while retaining every physics transition. +For locomotion data, ``--policy /path/to/exported/policy.pt`` accepts a feedforward TorchScript +policy consuming the task's ``policy`` observations. Match its observation layout, joint ordering +and action scale through task configuration. A policy trained in PhysX may need +``env.scene.robot.joint_ordering=physx`` when collecting with Newton. Multiple compatible +datasets can be passed together to ``train.py --dataset``. Use separate terrain seeds for final +evaluation; the learner's episode split alone does not establish terrain generalization. + +Collect and fit +~~~~~~~~~~~~~~~ + +Install the optional PhysicsNeMo Robotics dependency described above. These commands use +``--no-sync`` to preserve that separately installed experimental dependency: + +.. code-block:: bash + + uv run --no-sync python scripts/learning/nerd/record.py \ + --task Isaac-Cartpole-NeRD-Explore --output logs/nerd/cartpole/data.hdf5 \ + --num_envs 1024 --steps 1500 + uv run --no-sync python scripts/learning/nerd/train.py \ + --dataset logs/nerd/cartpole/data.hdf5 --output logs/nerd/cartpole/model \ + --wandb_project --wandb_entity + + uv run --no-sync python scripts/learning/nerd/record.py \ + --task Isaac-Velocity-Flat-Anymal-C-NeRD-Explore --output logs/nerd/anymal/data.hdf5 \ + --num_envs 4096 --steps 2400 --action_scale 1.5 --action_hold 4 + uv run --no-sync python scripts/learning/nerd/train.py \ + --dataset logs/nerd/anymal/data.hdf5 --output logs/nerd/anymal/model \ + --wandb_project --wandb_entity + +The default Cartpole recipe uses 30,000 updates and ten input frames. ANYmal C uses 60,000 +updates and thirteen frames, with its network geometry configured alongside the task. The +training loop samples history from flattened device-resident transitions, so memory does not +multiply by history length. Data must still fit on the selected device. ``--config`` accepts +JSON overrides of ``NeRDTrainCfg``; a supplied ``network`` dictionary replaces the task's +network geometry. For a smoke test, collect 32 steps in four worlds and fit +with ``--max_updates 3 --batch_size 16 --log_interval 1 --save_interval 1``. + +Collection writes native HDF5 episodes, a JSON representation descriptor, and the resolved +task YAML beside the dataset. Keep those files together. The recorder publishes preallocated +``state``, ``next_state``, ``control``, ``contacts``, ``terminated`` and ``truncated`` tensors. +Contact features come from the pre-state; applied effort is copied immediately before the solver; the terminal +successor is read before reset. An online consumer can read ``latest`` with +``record_to_episodes=False`` and feed replay or a prediction-error objective without a writer. + +Fitting splits complete episodes, repeats the first frame for reset prefixes, includes terminal +transitions, and rejects discontinuous episodes. It reports training/held-out normalized MSE +and held-out delta MSE. These measure dynamics prediction, not policy task performance. + +``last.pt`` contains network weights, fixed normalization, physical model identity, frame and +control conventions, optimizer state, sampler state and progress. Saves are atomic, happen at +least every 120 seconds between updates, and are uploaded synchronously when W&B is enabled. +SIGINT/SIGTERM finishes the current update and saves. Resume with: + +.. code-block:: bash + + uv run --no-sync python scripts/learning/nerd/train.py \ + --dataset logs/nerd/anymal/data.hdf5 --output logs/nerd/anymal/model \ + --checkpoint logs/nerd/anymal/model/last.pt --max_updates 80000 \ + --wandb_project --wandb_entity + +The checkpoint owns the normalization statistics. Data, split, network and optimization settings +must match; the update budget can increase without changing the learning-rate decay schedule. + +Model and precision options +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +``network_type`` selects the PhysicsNeMo network class. ParallelNeRD is the default. To use +VectorNeRD, supply a JSON configuration such as: + +.. code-block:: json + + { + "network_type": "physicsnemo_robotics.models.nerd:NewtonNeRDVectorTransformer", + "network": { + "num_heads": 8, "hidden_size": 192, "head_hidden": 64, + "temporal_depth": 4, "position_encoding": "relative_lag", + "temporal_attention_spans": [4, 3, 3, 3] + } + } + +``precision`` accepts ``float32`` (default), ``bfloat16`` or ``float16``. Parameters, normalized +targets and losses remain FP32. FP16 uses a saved gradient scaler and retries a skipped update +with the same samples; a nonfinite forward prediction stops training. ParallelNeRD FP16 requires +the dependency's attention matrix multiplications to disable autocast explicitly. The original +integration's runtime monkeypatch is not used. Development validation used PhysicsNeMo Robotics +branch ``feat/resident-dynamics-ownership`` (dependency commit ``e597e95``) on top of ``ee6d2e0``. It adds resident solver +integration, repeated-frame cache resets, canonical joint heading inputs and graph-safe stream bridging. +A released distribution of these APIs is not yet available; install the matching dependency source +provided with the integration review. Neither the original ``a81160ef`` revision nor stock PhysicsNeMo +alone supplies this interface. The PR remains dependent on publishing that compatible package. + +Launch distributed fitting through Torch's standard launcher: + +.. code-block:: bash + + uv run --no-sync python -m torch.distributed.run --standalone --nproc_per_node=2 \ + scripts/learning/nerd/train.py --dataset logs/nerd/anymal/data.hdf5 \ + --output logs/nerd/anymal/distributed --batch_size 4096 + +The batch size is global and must be divisible by the rank count. Each rank holds the dataset, +draws an independent sample stream and participates in gradient averaging. Validation is +partitioned across ranks. Rank zero owns logs, checkpoint publication and W&B uploads. +Checkpoints save every rank's sampler state; resume requires the same rank count. + +Train and evaluate a policy +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Select the learned task and pass its dynamics checkpoint through normal task overrides: + +.. code-block:: bash + + uv run --no-sync python scripts/reinforcement_learning/train.py \ + --rl_library rsl_rl --task Isaac-Cartpole-NeRD-Learned --num_envs 4096 \ + env.sim.physics.solver_cfg.checkpoint=logs/nerd/cartpole/model/last.pt + uv run --no-sync python scripts/reinforcement_learning/train.py \ + --rl_library rsl_rl --task Isaac-Velocity-Flat-Anymal-C-NeRD-Learned --num_envs 4096 \ + env.sim.physics.solver_cfg.checkpoint=logs/nerd/anymal/model/last.pt + + uv run --no-sync python scripts/reinforcement_learning/play.py \ + --rl_library rsl_rl --task Isaac-Velocity-Flat-Anymal-C-NeRD \ + --checkpoint /path/to/policy/model_299.pt --num_envs 16 --video --video_length 600 + +The last command evaluates the learned-physics policy in MuJoCo Warp. To evaluate the same +policy in learned physics, use ``-Learned`` and the same dynamics-checkpoint override. The +policy observation/action layouts and objectives agree between the two configurations. +ANYmal's acceleration penalty uses the control period in both solvers. Contact-force rewards, +torque penalties and physics randomizations absent from the model's inputs are disabled in both. + +Scope and migration +~~~~~~~~~~~~~~~~~~~ + +These are joint-backed rigid scenes with fixed physical parameters. The solver rejects a changed +physical model, unsupported timing, external body wrenches and solver-side motor drives. +Physical identity includes fixed-root orientation/height, body flags, collision filtering and +per-world topology. Horizontal reset-origin translations are ignored only for translation-invariant +flat scenes. Version-3 checkpoints serialize PhysicsNeMo's field normalizer, architecture, +gradient-scaler and per-rank sampler state. Existing version-2 raw applied-effort recordings can be +refit: they store world-space state, so the new heading encoder can derive its inputs from them. +Version-2 trained checkpoints are rejected explicitly. The canonical encoder zeros the mathematically +zero heading channel; using an old checkpoint's near-zero normalization scale could otherwise amplify +floating-point cancellation noise. Refit rather than silently changing that model's inputs. +Recollect datasets whose older physical fingerprint omitted the required properties. +Target-conditioned datasets and checkpoints are rejected. The former flat ANYmal ``-Effort-Explore`` +and ``-Effort-Learned`` tasks are now the standard ``-Explore`` and ``-Learned`` tasks, with no +compatibility aliases. Their raw applied-effort recordings retain the physical contract and can be refit to version 3. +Cartpole datasets and checkpoints collected with the former solver-side damping drive must be +recollected and refit for its actuator-owned damping. +ANYmal's flat-model fixed-plane probes do not represent rough terrain. Use the rough task family +to collect and fit the live-contact representation; flat and rough checkpoints are not interchangeable. + +Native HDF5 and checkpoints replace the experimental ``isaaclab_physicsnemo`` Zarr/artifact +launchers. Existing artifacts are not silently converted or treated as compatible. Historical +training results from that branch are not new validation of this workflow. The native integration +tests cover both tasks, terminal capture and retained values, offline update/resume, and repeated +learned-solver graph replay with asynchronous resets and intervening tensor allocations. +They also cover clipped actuator effort retained before buffer reuse, rejection of legacy target +models and solver-side drives, VectorNeRD, FP16 update/resume and two-rank CPU gradient +averaging against an independent combined-batch update. These are integration checks, not +long-horizon policy benchmarks or a multi-GPU scaling measurement. +The rough case also exercises live contact capture, contact-conditioned fitting and resume, +deployment with a different terrain and world count, and rejection of changed contact friction. + +Joint wrenches and contact forces are not learned outputs. Reconstruction requires the physical +model, contact law and acceleration in addition to position/velocity. Endpoint state alone does +not generally determine a unique contact-force distribution. Analytic reconstruction is being +validated separately and is not exposed as a learned-solver sensor capability. diff --git a/docs/source/experimental-features/supermimic.rst b/docs/source/experimental-features/supermimic.rst new file mode 100644 index 000000000000..0bba43e2a076 --- /dev/null +++ b/docs/source/experimental-features/supermimic.rst @@ -0,0 +1,217 @@ +.. _supermimic: + +SuperMimic motion tracking +========================== + +The experimental ``IsaacContrib-Motion-Tracking-G1`` task integrates the +`SuperMimic workflow `__ into Isaac Lab. +Use the standard reinforcement-learning entrypoints with ``--rl_library supermimic``. +The physical teacher runs Newton/MuJoCo Warp with CUDA graphs enabled. A learned world model +fits its recorded transitions, and the policy improves through differentiable imagined rollouts. + +This integration supports the MLP world model and the optional one-frame NeRD architecture. +It includes stationary padding of reference goals beyond recorded clip endpoints: padded poses +hold their position and orientation, with zero velocity. Recorded endpoint velocities remain unchanged. +The integration does not establish the paper's reported zero-failure result or an exact reproduction +of its motion selection. Evaluation reports failures over the actual prepared corpus. + +Ownership +--------- + +The task is an ordinary :class:`~isaaclab.envs.ManagerBasedRLEnv`. Its configuration lives in +``isaaclab_tasks.contrib.g1_tracking.g1_tracking_env_cfg`` and owns the scene, actuators, observations, +commands, rewards, terminations and recorder terms. The matching 29-DoF G1 asset variant and its +actuator settings are task-specific; substituting another G1 asset changes the dynamics. + +The remaining task modules have separate responsibilities: + +* ``mdp`` implements the native manager terms, including transition capture and adaptive reference sampling. +* ``tracking`` defines the state layout and the differentiable observation, reward and termination + expressions shared by real and imagined transitions. +* ``motion`` loads and resamples the reference corpus and constructs future goals. +* ``kinematics`` supplies the matching differentiable link geometry and joint order. + +``isaaclab_rl.supermimic`` owns the learner configuration, network composition, replay ring and +optimization. The native ``train_supermimic`` and ``play_supermimic`` entrypoint backends construct +the environment once and own stepping, checkpoint I/O, logging and evaluation. No external +SuperMimic package or launcher is required. The task does not import the learner, W&B or PhysicsNeMo. + +Both world-model choices use PhysicsNeMo Robotics' ``RunningNormalizer`` for world-model, policy and +critic statistics. Its updates use persistent device buffers without reading the count back to the +CPU. The G1 representation, reward and actor/critic update schedule remain native task/algorithm code. +Imagined rollouts call the network directly and retain action gradients through frozen world-model +weights; they do not use an inference session or detach state into a deployment cache. + +Transition capture and storage +------------------------------ + +``env.recorders.transition`` is disabled in the task's default configuration. SuperMimic training +enables ``TransitionRecorderCfg`` and reads its ``latest`` tensor dictionary after each control step. +This term captures the state, action, prior action, randomized context, reference index, reward, +episode step and termination flags. It captures the true successor before an automatic reset. +Its preallocated tensors are overwritten on the next step; a consumer that retains transitions +must copy them into its own storage. The learner does this in its time/world replay ring. + +The term uses the normal recorder lifecycle. It does not replace environment stepping or own replay. +By default, ``record_to_episodes=False`` avoids retaining per-world episodes. Offline consumers can +enable that setting and configure ``env.recorders.dataset_export_mode`` and the export directory +to use the same capture with the manager's batched episode storage and HDF5 exporter. +Other tasks can define their own :class:`~isaaclab.managers.RecorderTerm` schema at task-config level; +the core manager has no dependency on G1 or a learning algorithm. + +Prepare the robot and motions +----------------------------- + +Use a CUDA-capable Isaac Lab installation with the Newton backend and its standalone MJCF importer. +Install the compatible PhysicsNeMo Robotics ``feat/resident-dynamics-ownership`` (dependency commit ``e597e95``) dependency +alongside PhysicsNeMo ``2.1.1`` for the learner's shared normalization. Its released distribution is +still pending; the PR requires that separately supplied source revision. Task discovery and the +physical task remain independent of this learner dependency. Start with the 40-motion retargeted +LAFAN1 corpus. The preparation command downloads the robot +description and motions at a pinned dataset revision and records file hashes. The conversion +command is Isaac Lab's existing MJCF-to-USD tool: + +.. code-block:: bash + + uv run python scripts/tools/prepare_g1_tracking.py --output /data/g1_tracking + uv run python scripts/tools/convert_mjcf.py \ + /data/g1_tracking/g1_description/g1_29dof.xml \ + /data/g1_tracking/g1_usd/g1_29dof.usd + +The converter writes the composed asset to +``/data/g1_tracking/g1_usd/g1_29dof/g1_29dof.usda``. To reuse existing local inputs, supply +``--source_dir /path/to/prepared_inputs`` to the first command and choose an empty output directory. +That directory must contain ``g1_description/g1_29dof_rev_1_0.urdf``, its meshes and ``lafan1_g1``. +Local reuse records the local source and hashes, rather than claiming a downloaded revision. +The task defaults to the same layout under ``/data/g1_tracking``; the examples below +override those paths for an external data volume. + +For BONES-SEED, obtain the ``g1.tar.gz`` archive from the ``isaac/bones-seed-g1:1`` dataset and run: + +.. code-block:: bash + + uv run python scripts/tools/prepare_bones.py \ + --archive /data/bones-seed-g1/g1.tar.gz \ + --mjcf /data/g1_tracking/g1_description/g1_29dof.xml \ + --output /data/bones_prepared --workers 16 + +The converter reads named CSV columns, converts centimetres/degrees to metres/radians and resamples +120 Hz poses to 50 Hz. It writes retained poses under ``motions/``, per-clip decisions in +``motions.jsonl`` and filter settings and hashes in ``manifest.json``. The geometric and name filters +are explicit independent choices; the output is not claimed to match the authors' 129,514-clip subset. +Robot and motion data remain external inputs subject to their source licenses. BONES motion data +is attributed to Bones Studio. + +Train and resume +---------------- + +Train the MLP baseline through the unified entrypoint: + +.. code-block:: bash + + uv run python scripts/reinforcement_learning/train.py \ + --rl_library supermimic --task IsaacContrib-Motion-Tracking-G1 \ + --num_envs 4096 --log_dir logs/supermimic/g1_mlp \ + env.scene.robot.spawn.usd_path=/data/g1_tracking/g1_usd/g1_29dof/g1_29dof.usda \ + env.commands.motion.kinematics_path=/data/g1_tracking/g1_description/g1_29dof.xml \ + env.commands.motion.motion_dir=/data/g1_tracking/lafan1_g1 + +The default agent uses BF16 neural projections, FP32 geometry and compiled tensor rollouts. +Task settings use ``env.*`` overrides; optimizer, replay and runner settings use ``agent.*``. +The analytical imagined objective requires the task's observation and termination functions. +Observation history, stateless policy noise and task reward parameters remain configurable. +The runner rejects incompatible observation functions or layouts instead of silently optimizing +a different imagined task. + +For a short integration smoke test, use ``--num_envs 32 --max_iterations 2`` and append: + +.. code-block:: bash + + agent.buffer_steps=16 agent.warmup_steps=4 agent.world_model_window=2 \ + agent.world_model_batch_size=32 agent.world_model_batches=1 agent.policy_batch_size=32 \ + agent.horizon=4 agent.critic_epochs=1 agent.critic_minibatches=1 \ + agent.world_model_only_iterations=0 agent.log_interval=1 \ + env.observations.policy.history_length=3 env.commands.motion.goal_frames=2 + +These reduced settings test the pipeline; they are not a learning benchmark. For the prepared +BONES corpus, replace the motion directory with ``/data/bones_prepared/motions`` and set +``env.commands.motion.source_fps=50``. For a nine-hour budget, add +``--max_iterations 250000 agent.max_hours=9``. The budget is cumulative when resuming and is +checked after each update. Initial environment setup and replay warmup are outside that budget. + +Optional NeRD world model +~~~~~~~~~~~~~~~~~~~~~~~~~ + +Set ``agent.world_model=nerd-one-frame`` to select PhysicsNeMo Robotics' ``ParallelNeRD`` with +four attention spans of one. The policy, critic, task, replay and optimization objective are shared +with the MLP configuration. NeRD is used inside the learner; the physical teacher still uses +MuJoCo Warp. This option does not deploy a learned collision or constraint solver into Newton. + +This experimental option requires a compatible PhysicsNeMo Robotics installation exposing +``physicsnemo_robotics.models.parallel_nerd.ParallelNeRD``. The NeRD network is imported only when selected; +both architectures use the shared PhysicsNeMo normalizer. The integration uses PhysicsNeMo Robotics +``0.1.0.dev0`` from the dependency branch above and PhysicsNeMo ``2.1.1`` in the development environment. +A clean, released package combination is not established here; +the experimental dependency must be supplied separately. + +The offline NeRD workflow below uses that same network implementation with temporal attention. +Its generalized-coordinate state and fixed offline normalization differ from SuperMimic's +task-specific state and online statistics. Both normalization implementations belong to PhysicsNeMo; +the application chooses the convention its algorithm needs. Existing SuperMimic normalization buffer +names and checkpoint state remain compatible. + +Logging and preemption +~~~~~~~~~~~~~~~~~~~~~~ + +Every run writes resolved configurations, a run manifest and ``train_log.csv``. To enable W&B, +authenticate its SDK and append ``agent.wandb_project= agent.wandb_entity=``. +An optional stable ``agent.wandb_run_id=`` lets resumed runs continue the same history. +The runner records the actual run identity in ``wandb-run.json`` and the checkpoint. + +``last.pt`` contains the world model, policy, critic, their optimizers, normalizers, adaptive sampler +statistics and semantic configurations. It is replaced atomically after the first update, at the +configured interval, after 120 seconds by default, and on completion. When W&B is enabled, each +save is uploaded as a version of ``run--checkpoint`` with the ``latest`` alias; the runner waits +for upload before reporting ``durable_checkpoint_iteration``. Milestones also save +``policy_.pt`` for playback. SIGINT/SIGTERM request a save after the current update. +An immediate kill can lose progress since the last completed checkpoint. + +To resume, repeat the original task and learning overrides, add ``--checkpoint /path/to/last.pt`` +and increase ``--max_iterations``. Logging, save intervals and budget settings can change. +Physical state and replay are rebuilt; model and optimizer state and sampler statistics resume. +This is a training continuation, not a bitwise continuation of the simulator trajectory. +The runner accepts its native checkpoint metadata; legacy external-run checkpoints require +explicit conversion before use. + +Evaluate and record video +------------------------- + +Playback restores the native checkpoint's task settings and evaluates every clip exactly once: + +.. code-block:: bash + + uv run python scripts/reinforcement_learning/play.py \ + --rl_library supermimic --checkpoint logs/supermimic/g1_mlp/policy_014000.pt \ + --num_envs 64 --output logs/supermimic/g1_mlp/evaluation.json + +The output includes the failure count, failed names, per-clip returns and lengths, and both +clip-weighted and frame-weighted normalized returns. Failed tails remain in the full-clip +denominator. Use these physical evaluation metrics to assess policy quality; dynamics loss +measures prediction fit and can improve without improving tracking success. + +Use ``--motion_dir`` to select another corpus in the same pose schema and source frame rate. +``--asset_path`` and ``--kinematics_path`` relocate the matching assets. Add the standard +``--video --video_length 600`` options to record playback with Isaac Lab's native renderer. +For an interpretable video, select a small motion directory and a small ``--num_envs`` count. + +Validation scope +---------------- + +The native tests cover terminal capture, observation/kinematics agreement, episode export, +stationary goal padding, replay boundaries, complete evaluation batches and gradients through +compiled MLP/NeRD dynamics. Asset-dependent tests use ``ISAACLAB_G1_TRACKING_DATA`` pointing to +the prepared directory above. The dependency-direction tests prevent reintroducing an external +training forwarder or moving learner ownership into the task. + +The source adaptation preserves the upstream SuperMimic MIT notice in +``docs/licenses/dependencies/supermimic-license.txt``. diff --git a/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py b/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py index 1659510f7412..5a232f90fbc2 100644 --- a/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py +++ b/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py @@ -515,7 +515,6 @@ def annotate_episode_in_manual_mode( print(f'\tThe episode will be replayed again for re-marking subtask signals for the eef "{eef_name}".\n') - annotated_episode = env.recorder_manager.get_episode(0) for ( subtask_term_signal_name, subtask_term_signal_action_index, @@ -523,7 +522,9 @@ def annotate_episode_in_manual_mode( # subtask termination signal is false until subtask is complete, and true afterwards subtask_signals = torch.ones(len(episode.data["actions"]), dtype=torch.bool) subtask_signals[:subtask_term_signal_action_index] = False - annotated_episode.add(f"obs/datagen_info/subtask_term_signals/{subtask_term_signal_name}", subtask_signals) + env.recorder_manager.add_to_episodes( + f"obs/datagen_info/subtask_term_signals/{subtask_term_signal_name}", subtask_signals.unsqueeze(0), [0] + ) if args_cli.annotate_subtask_start_signals: for ( @@ -532,8 +533,8 @@ def annotate_episode_in_manual_mode( ) in subtask_start_signal_action_indices.items(): subtask_signals = torch.ones(len(episode.data["actions"]), dtype=torch.bool) subtask_signals[:subtask_start_signal_action_index] = False - annotated_episode.add( - f"obs/datagen_info/subtask_start_signals/{subtask_start_signal_name}", subtask_signals + env.recorder_manager.add_to_episodes( + f"obs/datagen_info/subtask_start_signals/{subtask_start_signal_name}", subtask_signals.unsqueeze(0), [0] ) return True diff --git a/scripts/learning/nerd/record.py b/scripts/learning/nerd/record.py new file mode 100644 index 000000000000..a4ce5a08db47 --- /dev/null +++ b/scripts/learning/nerd/record.py @@ -0,0 +1,89 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Collect a registered task's native NeRD transitions into an HDF5 dataset.""" + +import argparse +import json +from pathlib import Path + +import gymnasium as gym +import torch + +from isaaclab.app import add_launcher_args, launch_simulation +from isaaclab.managers import DatasetExportMode +from isaaclab.utils.assets import read_file +from isaaclab.utils.io import dump_yaml + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.utils import resolve_task_config, setup_preset_cli + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--task", required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--num_envs", type=int, default=1024) + parser.add_argument("--steps", type=int, default=1500) + parser.add_argument("--action_scale", type=float, default=1.5, help="Uniform normalized action bound.") + parser.add_argument( + "--policy", help="Exported feedforward TorchScript policy path or URL, consuming policy observations." + ) + parser.add_argument("--action_hold", type=int, default=1, help="Environment steps between action updates.") + parser.add_argument("--seed", type=int, default=42) + add_launcher_args(parser) + args, overrides = setup_preset_cli(parser) + if min(args.num_envs, args.steps, args.action_hold) < 1 or args.action_scale <= 0: + parser.error("World/step counts, action hold and action scale must be positive.") + path = args.output.resolve() + if path.suffix != ".hdf5": + parser.error("--output must end in .hdf5.") + if path.exists() or path.with_suffix(".json").exists(): + parser.error("The output dataset already exists; choose another path.") + cfg, _ = resolve_task_config(args.task, "", overrides=overrides) + if cfg.recorders.transition is None: + parser.error("Choose a task with a configured NeRD transition recorder, such as a -Explore task.") + cfg.scene.num_envs = args.num_envs + cfg.sim.device = args.device + cfg.seed = args.seed + cfg.recorders.transition.record_to_episodes = True + cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_ALL + cfg.recorders.dataset_export_dir_path = str(path.parent) + cfg.recorders.dataset_filename = path.stem + cfg.recorders.export_in_close = True + path.parent.mkdir(parents=True, exist_ok=True) + dump_yaml(str(path.with_suffix(".yaml")), cfg) + torch.manual_seed(args.seed) + with launch_simulation(cfg, args): + env = gym.make(args.task, cfg=cfg).unwrapped + try: + observations, _ = env.reset() + policy = torch.jit.load(read_file(args.policy), map_location=env.device).eval() if args.policy else None + capture = env.recorder_manager.get_term("transition") + metadata = { + "task": args.task, + "seed": args.seed, + "action_scale": args.action_scale, + "policy": args.policy, + "action_hold": args.action_hold, + "dynamics": capture.dynamics.descriptor, + } + path.with_suffix(".json").write_text(json.dumps(metadata, indent=2) + "\n") + with torch.inference_mode(): + actions = torch.empty(env.num_envs, env.action_manager.total_action_dim, device=env.device) + for step in range(args.steps): + if step % args.action_hold == 0: + if policy is None: + actions.uniform_(-args.action_scale, args.action_scale) + else: + actions.copy_(policy(observations["policy"])) + observations, _, _, _, _ = env.step(actions) + finally: + env.close() + print(f"Recorded native transitions to {path}") + + +if __name__ == "__main__": + main() diff --git a/scripts/learning/nerd/train.py b/scripts/learning/nerd/train.py new file mode 100644 index 000000000000..e0efff01a8e3 --- /dev/null +++ b/scripts/learning/nerd/train.py @@ -0,0 +1,161 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Fit temporal NeRD to native HDF5 episodes, with resumable checkpoints and optional W&B uploads.""" + +import argparse +import json +import os +import signal +import time +from contextlib import nullcontext +from pathlib import Path + +import torch +import torch.distributed as dist + +from isaaclab.utils import to_dict + +from isaaclab_rl.nerd import NeRDTrainCfg, OfflineNeRD + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.utils import load_cfg_from_registry + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dataset", type=Path, nargs="+", required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--config", type=Path, help="JSON overrides of NeRDTrainCfg, including network geometry.") + parser.add_argument("--max_updates", type=int) + parser.add_argument("--batch_size", type=int) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--checkpoint", type=Path, help="Resume a native checkpoint.") + parser.add_argument("--save_interval", type=int, default=1000) + parser.add_argument("--save_seconds", type=float, default=120.0) + parser.add_argument("--log_interval", type=int, default=100) + parser.add_argument("--wandb_project") + parser.add_argument("--wandb_entity") + parser.add_argument("--wandb_name") + args = parser.parse_args() + if min(args.save_interval, args.save_seconds, args.log_interval) <= 0: + parser.error("Checkpoint and logging intervals must be positive.") + distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1 + if distributed: + if args.device.startswith("cuda"): + args.device = f"cuda:{os.environ['LOCAL_RANK']}" + torch.cuda.set_device(args.device) + dist.init_process_group(backend="nccl" if args.device.startswith("cuda") else "gloo") + rank = dist.get_rank() if distributed else 0 + task = json.loads(args.dataset[0].with_suffix(".json").read_text())["task"] + task_cfg = load_cfg_from_registry(task, "nerd_cfg_entry_point") + settings = task_cfg if isinstance(task_cfg, dict) else to_dict(task_cfg) + if args.config: + settings.update(json.loads(args.config.read_text())) + cfg = NeRDTrainCfg(**settings) + if args.max_updates is not None: + cfg.max_updates = args.max_updates + if args.batch_size is not None: + cfg.batch_size = args.batch_size + args.output.mkdir(parents=True, exist_ok=True) + target = args.output / "last.pt" + if target.exists() and not args.checkpoint: + parser.error("Output contains last.pt; use --checkpoint to resume or choose another output.") + saved = torch.load(args.checkpoint, map_location="cpu", weights_only=True) if args.checkpoint else None + settings = to_dict(cfg) + if saved: + previous = saved["training_cfg"] + if any(previous[key] != value for key, value in settings.items() if key not in ("max_updates", "compile")): + parser.error("Resume requires unchanged optimization settings, split and model geometry.") + learner = OfflineNeRD(args.dataset, cfg, args.device, checkpoint=saved) + if learner.updates >= cfg.max_updates: + parser.error("The requested update limit has already been reached.") + logging = nullcontext(None) + if args.wandb_project and rank == 0: + import wandb + + identity = saved.get("wandb") if saved else None + resume_identity = ( + identity + and identity["project"] == args.wandb_project + and (args.wandb_entity is None or identity["entity"] == args.wandb_entity) + ) + logging = wandb.init( + project=args.wandb_project, + entity=identity["entity"] if resume_identity else args.wandb_entity, + name=args.wandb_name, + id=identity["id"] if resume_identity else None, + resume="allow", + config=settings, + allow_val_change=True, + dir=str(args.output), + ) + stopped = False + + def stop(signum, frame): + nonlocal stopped + stopped = True + + handlers = {signum: signal.signal(signum, stop) for signum in (signal.SIGINT, signal.SIGTERM)} + try: + metrics_file = (args.output / "metrics.jsonl").open("a") if rank == 0 else nullcontext(None) + with logging as run, metrics_file as metrics: + if run: + run.define_metric("update") + run.define_metric("*", step_metric="update") + last_save = time.monotonic() + while learner.updates < cfg.max_updates: + values = learner.update() + checkpoint_due = ( + learner.updates % args.save_interval == 0 or time.monotonic() - last_save >= args.save_seconds + ) + finished = learner.updates == cfg.max_updates or stopped + if distributed: + decisions = torch.tensor([checkpoint_due, finished], dtype=torch.int32, device=args.device) + dist.all_reduce(decisions, op=dist.ReduceOp.MAX) + checkpoint_due, finished = decisions.tolist() + if learner.updates % args.log_interval == 0 or finished or checkpoint_due: + values |= learner.validate() + values["update"] = learner.updates + if metrics: + metrics.write(json.dumps(values) + "\n") + metrics.flush() + print(json.dumps(values), flush=True) + if run: + run.log(values) + if checkpoint_due or finished: + checkpoint = learner.checkpoint() + if rank == 0: + checkpoint["training_cfg"] = settings + checkpoint["wandb"] = ( + {"id": run.id, "project": run.project, "entity": run.entity} if run else None + ) + temporary = target.with_suffix(".tmp") + with temporary.open("wb") as stream: + torch.save(checkpoint, stream) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, target) + if run: + artifact = wandb.Artifact(f"nerd-{run.id}", type="model", metadata={"update": learner.updates}) + artifact.add_file(str(target), name="last.pt") + logged = run.log_artifact(artifact, aliases=["latest", f"update-{learner.updates}"]) + if run.settings.mode != "offline": + logged.wait(timeout=180) + run.log({"update": learner.updates, "checkpoint/durable_update": learner.updates}) + last_save = time.monotonic() + if finished: + break + finally: + for signum, handler in handlers.items(): + signal.signal(signum, handler) + + +if __name__ == "__main__": + try: + main() + finally: + if dist.is_initialized(): + dist.destroy_process_group() diff --git a/scripts/tools/prepare_bones.py b/scripts/tools/prepare_bones.py new file mode 100644 index 000000000000..ff04d30d1004 --- /dev/null +++ b/scripts/tools/prepare_bones.py @@ -0,0 +1,247 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""Convert the public BONES-SEED archive into metric G1 poses and a versioned selection manifest. + +The name exclusions adapt SONIC's published filter, retaining floor acrobatics. The geometric +checks follow SuperMimic Appendix E; the unreported stationary gates are explicit below. +This independent selection is not claimed to match the authors' exact 129514-file manifest. +""" + +import argparse +import csv +import hashlib +import io +import json +import multiprocessing as mp +import tarfile +from collections import Counter, deque +from pathlib import Path + +import mujoco +import numpy as np +import torch +from scipy.spatial.transform import Rotation + +from isaaclab_tasks.contrib.g1_tracking.kinematics import G1_JOINT_NAMES +from isaaclab_tasks.contrib.g1_tracking.motion import CONTROL_DT, resample_motion + +# Derived from NVlabs/GR00T-WholeBodyControl filter_and_copy_bones_data.py, +# blob f09061012a1f1d1da3ef35d614440629138659cd. Floor acrobatics remain included. +EXCLUDED_NAMES = ( + "bed", + "bike", + "chair", + "climb", + "com_up_50cm", + "sitting", + "step_on", + "seat", + "table", + "sit_", + "ladder", + "crutch", + "_ride_", + "scooter", + "stepdown", + "box_hspu", + "50cm_box_", + "on_box", + "fall_from", + "on_1m", + "form_box", + "off_1m", + "230m", + "jump_over_obstacle_", + "lift_crate_come_up_", + "stair", + "box_jump", + "monkey_jump", + "box_dips", + "walking_on_edge", + "push_obstacle", + "jump_off_50cm", + "look_over_fence", +) +FILTER = { + "elevated_height_m": 0.20, + "elevated_seconds": 0.6, + "penetration_m": -0.12, + "sustained_penetration_m": -0.08, + "penetration_fraction": 0.8, + "floating_clearance_m": 0.02, + "floating_seconds": 0.5, + "standing_pelvis_m": 0.55, + "upright_cosine": 0.9, + "stationary_max_link_speed_m_s": 0.2, + "stationary_gate_provenance": "Independent choices; these gates are not specified numerically in the paper.", +} + + +def initialize_worker(mjcf_path, output, source_fps): + """Own one MuJoCo model/data pair per conversion process; never share mutable simulator data.""" + global MODEL, DATA, FLOOR, FEET, OUTPUT, SOURCE_FPS + torch.set_num_threads(1) + spec = mujoco.MjSpec.from_file(mjcf_path) + spec.worldbody.add_geom(name="filter_floor", type=mujoco.mjtGeom.mjGEOM_PLANE, size=[0, 0, 0.1]) + MODEL = spec.compile() + DATA = mujoco.MjData(MODEL) + FLOOR = MODEL.geom("filter_floor").id + FEET = [] + for side in ("left", "right"): + body = MODEL.body(f"{side}_ankle_roll_link").id + geoms = np.flatnonzero((MODEL.geom_bodyid == body) & ((MODEL.geom_contype | MODEL.geom_conaffinity) != 0)) + if len(geoms) == 0: + raise ValueError(f"No collision geometry on {side} foot.") + FEET.append(geoms) + joint_names = [MODEL.joint(j).name for j in range(MODEL.njnt) if MODEL.jnt_type[j] != mujoco.mjtJoint.mjJNT_FREE] + if joint_names != G1_JOINT_NAMES: + raise ValueError("Filtering model joint order differs from the pose schema.") + OUTPUT, SOURCE_FPS = Path(output), source_fps + + +def convert_csv(data): + """BONES header-defined cm/degree/Euler poses to metric xyzw/radian poses in task joint order.""" + header = next(csv.reader(io.StringIO(data.decode("utf-8-sig")))) + expected = [ + "Frame", + *[f"root_translate{x}" for x in "XYZ"], + *[f"root_rotate{x}" for x in "XYZ"], + *[f"{name}_dof" for name in G1_JOINT_NAMES], + ] + if len(set(header)) != len(header) or set(header) != set(expected): + raise ValueError(f"Unexpected BONES CSV fields: {header}") + raw = np.loadtxt(io.BytesIO(data), delimiter=",", skiprows=1, ndmin=2) + raw = raw[:, [header.index(name) for name in expected]] + if len(raw) < 4 or not np.isfinite(raw).all() or not np.allclose(np.diff(raw[:, 0]), 1, atol=0, rtol=0): + raise ValueError("Invalid pose values or nonconsecutive source frame indices.") + quat = Rotation.from_euler("xyz", raw[:, 4:7], degrees=True).as_quat() + return np.concatenate((raw[:, 1:4] / 100, quat, np.deg2rad(raw[:, 7:])), axis=1).astype(np.float32) + + +def sustained(mask, seconds): + n = int(round(seconds / CONTROL_DT)) + return len(mask) >= n and bool((np.convolve(mask.astype(np.int32), np.ones(n, dtype=np.int32), "valid") == n).any()) + + +def process_motion(name, data): + record = {"source": name, "source_sha256": hashlib.sha256(data).hexdigest()} + raw = convert_csv(data) + p, quat, q = resample_motion(torch.from_numpy(raw), SOURCE_FPS) + poses = torch.cat((p, quat, q), dim=-1).numpy() + n = len(poses) + zmin, clearance, positions = np.empty(n), np.empty(n), [] + for i, pose in enumerate(poses): + DATA.qpos[:] = np.concatenate((pose[:3], pose[[6, 3, 4, 5]], pose[7:])) + mujoco.mj_kinematics(MODEL, DATA) + zmin[i] = DATA.xpos[1:, 2].min() + clearance[i] = min(mujoco.mj_geomDistance(MODEL, DATA, int(g), FLOOR, 10, None) for gs in FEET for g in gs) + positions.append(DATA.xpos[1:].copy()) + max_speed = np.linalg.norm(np.gradient(np.stack(positions), CONTROL_DT, axis=0), axis=-1).max(axis=-1) + up = Rotation.from_quat(poses[:, 3:7]).as_matrix()[:, 2, 2] + floating = ( + (clearance > FILTER["floating_clearance_m"]) + & (poses[:, 2] > FILTER["standing_pelvis_m"]) + & (up > FILTER["upright_cosine"]) + & (max_speed < FILTER["stationary_max_link_speed_m_s"]) + ) + reason = None + if sustained(zmin > FILTER["elevated_height_m"], FILTER["elevated_seconds"]): + reason = "elevated_support" + elif (zmin < FILTER["penetration_m"]).any() or (zmin < FILTER["sustained_penetration_m"]).mean() > FILTER[ + "penetration_fraction" + ]: + reason = "ground_penetration" + elif sustained(floating, FILTER["floating_seconds"]): + reason = "floating_feet" + record.update(source_frames=len(raw), frames=n, reason=reason) + if reason is None: + path = Path(name).with_suffix(".npy") + target = OUTPUT / "motions" / path + target.parent.mkdir(parents=True, exist_ok=True) + np.save(target, poses, allow_pickle=False) + record.update(path=path.as_posix(), pose_sha256=hashlib.sha256(poses.tobytes()).hexdigest()) + return record + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--archive", required=True) + parser.add_argument("--mjcf", required=True, help="Matching 29-DoF G1 MJCF from prepare_g1_tracking.py.") + parser.add_argument("--output", required=True) + parser.add_argument("--workers", type=int, default=16) + parser.add_argument("--source_fps", type=float, default=120) + args = parser.parse_args() + output = Path(args.output) + if output.exists() and any(output.iterdir()): + parser.error("Output must be empty; do not mix dataset versions.") + output.mkdir(parents=True, exist_ok=True) + counts, kept_frames, total = Counter(), 0, 0 + manifest_path = output / "motions.jsonl" + with ( + manifest_path.open("w") as log, + mp.get_context("spawn").Pool( + args.workers, initializer=initialize_worker, initargs=(args.mjcf, str(output), args.source_fps) + ) as pool, + ): + pending = deque() + + def collect_one(): + nonlocal kept_frames + result = pending.popleft().get() + reason = result["reason"] or "kept" + counts[reason] += 1 + if reason == "kept": + kept_frames += result["frames"] + log.write(json.dumps(result) + "\n") + if sum(counts.values()) % 5000 == 0: + log.flush() + print("PROGRESS", dict(counts), "kept_frames", kept_frames, flush=True) + + with tarfile.open(args.archive, "r|gz") as archive: + for member in archive: + if not member.isfile() or not member.name.endswith(".csv"): + continue + name = Path(member.name) + if name.is_absolute() or ".." in name.parts: + raise ValueError(f"Unsafe archive path: {name}") + total += 1 + keyword = next((key for key in EXCLUDED_NAMES if key in name.stem.lower()), None) + if keyword: + counts["name"] += 1 + log.write(json.dumps({"source": member.name, "reason": "name", "keyword": keyword}) + "\n") + continue + pending.append(pool.apply_async(process_motion, (member.name, archive.extractfile(member).read()))) + if len(pending) >= 2 * args.workers: + collect_one() + while pending: + collect_one() + summary = { + "source": "isaac/bones-seed-g1:1", + "attribution": "Motion Data by Bones Studio", + "raw_clips": total, + "counts": dict(counts), + "kept_frames": kept_frames, + "kept_hours": (kept_frames - counts["kept"]) * CONTROL_DT / 3600, + "source_fps": args.source_fps, + "prepared_fps": 1 / CONTROL_DT, + "layout": "root_xyz_m,root_quaternion_xyzw,29_joints_rad", + "joint_names": G1_JOINT_NAMES, + "name_exclusions": EXCLUDED_NAMES, + "geometric_filter": FILTER, + "exact_paper_subset": False, + "selection_sha256": hashlib.sha256(manifest_path.read_bytes()).hexdigest(), + } + (output / "manifest.json").write_text(json.dumps(summary, indent=2) + "\n") + print("PREPARATION_COMPLETE", json.dumps(summary), flush=True) + if not total or not counts["kept"]: + raise ValueError("No motions retained; inspect the selection manifest.") + + +if __name__ == "__main__": + main() diff --git a/scripts/tools/prepare_g1_tracking.py b/scripts/tools/prepare_g1_tracking.py new file mode 100644 index 000000000000..7aabd90b680e --- /dev/null +++ b/scripts/tools/prepare_g1_tracking.py @@ -0,0 +1,96 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Prepare pinned G1 tracking inputs; convert the resulting MJCF with scripts/tools/convert_mjcf.py. + +Downloads the retargeted LAFAN1 motions and matching Unitree robot description, retaining the source +revision in a manifest. The datasets retain their original licenses; this script does not redistribute them. +""" + +import argparse +import hashlib +import json +import shutil +import tempfile +import urllib.request +import xml.etree.ElementTree as ET +from pathlib import Path + +import mujoco + +REVISION = "ce1572906efe6157840e8474d5a0d7aa87481e74" +REPOSITORY = "lvhaidong/LAFAN1_Retargeting_Dataset" + + +def main() -> None: + """Download the inputs or reuse a local copy, then compile a free-floating 29-DoF G1 MJCF.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--source_dir", type=Path, help="Existing prepared directory with g1_description and lafan1_g1." + ) + args = parser.parse_args() + output = args.output.resolve() + if output.exists() and any(output.iterdir()): + parser.error("Output must be empty to keep preparation reproducible.") + output.mkdir(parents=True, exist_ok=True) + if args.source_dir: + source = {"local_directory": str(args.source_dir.resolve())} + for name in ("g1_description", "lafan1_g1"): + shutil.copytree(args.source_dir / name, output / name) + else: + source = {"repository": REPOSITORY, "revision": REVISION} + paths = ["robot_description/g1/g1_29dof_rev_1_0.urdf"] + for directory in ("g1", "robot_description/g1/meshes"): + with urllib.request.urlopen( + f"https://huggingface.co/api/datasets/{REPOSITORY}/tree/{REVISION}/{directory}?limit=1000" + ) as response: + paths.extend(entry["path"] for entry in json.load(response) if entry["type"] == "file") + for name in paths: + relative = Path(name) + if relative.parts[0] == "g1": + target = output / "lafan1_g1" / relative.name + else: + target = output / "g1_description" / relative.relative_to("robot_description/g1") + target.parent.mkdir(parents=True, exist_ok=True) + urllib.request.urlretrieve( + f"https://huggingface.co/datasets/{REPOSITORY}/resolve/{REVISION}/{name}", target + ) + description = output / "g1_description" + robot = ET.parse(description / "g1_29dof_rev_1_0.urdf") + root = robot.getroot() + for element in root.findall("mujoco"): + root.remove(element) + compiler = ET.SubElement(ET.SubElement(root, "mujoco"), "compiler") + compiler.attrib.update( + meshdir=".", balanceinertia="true", discardvisual="false", fusestatic="false", strippath="false" + ) + with tempfile.NamedTemporaryFile(dir=description, suffix=".urdf") as temporary: + robot.write(temporary.name) + model = mujoco.MjModel.from_xml_path(temporary.name) + mjcf = description / "g1_29dof.xml" + mujoco.mj_saveLastXML(str(mjcf), model) + compiled = ET.parse(mjcf) + pelvis = compiled.find(".//body[@name='pelvis']") + if pelvis is None: + raise ValueError("The pinned robot must retain its pelvis body.") + pelvis.insert(0, ET.Element("freejoint", name="floating_base_joint")) + compiled.write(mjcf) + model = mujoco.MjModel.from_xml_path(str(mjcf)) + if model.nq != 36 or model.nv != 35: + raise ValueError("Expected a free-floating 29-DoF robot.") + source["sha256"] = {} + for path in sorted(output.rglob("*")): + if path.is_file(): + with path.open("rb") as stream: + source["sha256"][path.relative_to(output).as_posix()] = hashlib.file_digest( + stream, "sha256" + ).hexdigest() + (output / "source.json").write_text(json.dumps(source, indent=2) + "\n") + print(f"Prepared {mjcf}; convert it with scripts/tools/convert_mjcf.py before training.") + + +if __name__ == "__main__": + main() diff --git a/scripts/tools/test/test_prepare_bones.py b/scripts/tools/test/test_prepare_bones.py new file mode 100644 index 000000000000..8692dbb50059 --- /dev/null +++ b/scripts/tools/test/test_prepare_bones.py @@ -0,0 +1,54 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Independent unit, frame-rate and physical-support contracts for BONES preparation.""" + +import io +import os +from pathlib import Path + +import numpy as np +import pytest +from scipy.spatial.transform import Rotation + +from isaaclab_tasks.contrib.g1_tracking.kinematics import G1_JOINT_NAMES + +from scripts.tools import prepare_bones + + +@pytest.fixture +def elevated_clip(): + columns = [ + "Frame", + *[f"root_translate{x}" for x in "XYZ"], + *[f"root_rotate{x}" for x in "XYZ"], + *[f"{name}_dof" for name in G1_JOINT_NAMES], + ] + values = np.zeros((121, len(columns))) + values[:, 0], values[:, 3], values[:, 6], values[:, 7] = np.arange(121), 300, 90, 45 + order = np.arange(len(columns))[::-1] + stream = io.StringIO() + np.savetxt(stream, values[:, order], delimiter=",", header=",".join(columns[i] for i in order), comments="") + return stream.getvalue().encode() + + +def test_conversion_uses_header_order_and_named_units(elevated_clip): + poses = prepare_bones.convert_csv(elevated_clip) + np.testing.assert_allclose(poses[:, :3], np.tile([0, 0, 3], (121, 1))) + np.testing.assert_allclose(poses[:, 7], np.pi / 4) + np.testing.assert_allclose(Rotation.from_quat(poses[0, 3:7]).apply([1, 0, 0]), [0, 1, 0], atol=1e-6) + with pytest.raises(ValueError, match="fields"): + prepare_bones.convert_csv(elevated_clip.replace(b"root_translateX", b"root_translateY")) + + +def test_filter_resamples_and_rejects_missing_support(elevated_clip, tmp_path): + data_dir = os.environ.get("ISAACLAB_G1_TRACKING_DATA") + if not data_dir: + pytest.skip("Set ISAACLAB_G1_TRACKING_DATA to the prepared G1 tracking directory.") + mjcf = Path(data_dir) / "g1_description/g1_29dof.xml" + prepare_bones.initialize_worker(str(mjcf), str(tmp_path), 120) + result = prepare_bones.process_motion("session/elevated.csv", elevated_clip) + assert result["reason"] == "elevated_support" and result["frames"] == 51 + assert not list(tmp_path.rglob("*.npy")) diff --git a/source/isaaclab/changelog.d/quaternion-zero-gradient.fixed.rst b/source/isaaclab/changelog.d/quaternion-zero-gradient.fixed.rst new file mode 100644 index 000000000000..4ab770ab120a --- /dev/null +++ b/source/isaaclab/changelog.d/quaternion-zero-gradient.fixed.rst @@ -0,0 +1,2 @@ +* Fixed NaN gradients in ``axis_angle_from_quat`` at identity rotations by keeping the + unselected exact-form division finite during the Taylor approximation. diff --git a/source/isaaclab/changelog.d/recorder-capture.changed.rst b/source/isaaclab/changelog.d/recorder-capture.changed.rst new file mode 100644 index 000000000000..3e2c6396515e --- /dev/null +++ b/source/isaaclab/changelog.d/recorder-capture.changed.rst @@ -0,0 +1,21 @@ +* Replaced per-world recorder episode containers with batched field allocations and + independent insertion indices, reusing storage across resets and growing without dropping samples. +* Added :meth:`~isaaclab.managers.RecorderManager.get_term` for direct access to captured batches. +* Preserve slices through recorder reset callbacks; terms accept slices or indices directly. +* **Breaking:** Removed the implicit observation computation before RL post-step recorder callbacks. + Physical capture no longer consumes extra observation noise. Custom terms needing observations + of the completed transition must call ``env.observation_manager.compute_group(name, update_history=False)`` + explicitly; ``env.obs_buf`` still contains the preceding policy input. Use a noiseless group without + history or delay for current measurements that do not consume policy noise draws. +* **Breaking:** Changed :meth:`~isaaclab.managers.RecorderManager.get_episode` to return an owned + snapshot with tensor leaves. To update live recording, use ``add_to_episodes`` and + ``set_success_to_episodes``; do not mutate the snapshot. Set ``buffer_initial_capacity`` + to control the initial allocation per field. +* **Breaking:** Removed ``EpisodeData.add()`` and ``pre_export()``. Collection now belongs to + ``RecorderManager.add_to_episodes()``; assign nested tensors with leading sample dimensions + to ``EpisodeData.data`` when constructing an episode directly. Playback and HDF5 loading remain supported. +* **Breaking:** Removed the unused ``RecorderTerm.close(file_path)`` hook. Dataset handlers own + file finalization; applications must explicitly close any additional resources they own. + +* Documented independent capture, retention and export choices, with task-selected schemas and + consumer-owned processing. Core recorder storage remains independent of task and learning packages. diff --git a/source/isaaclab/changelog.d/recorder-capture.major b/source/isaaclab/changelog.d/recorder-capture.major new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab/changelog.d/recorder-inference-lifetime.fixed.rst b/source/isaaclab/changelog.d/recorder-inference-lifetime.fixed.rst new file mode 100644 index 000000000000..1302fe5c27d3 --- /dev/null +++ b/source/isaaclab/changelog.d/recorder-inference-lifetime.fixed.rst @@ -0,0 +1,2 @@ +* Kept retained recorder allocations writable after leaving PyTorch inference mode, including + grown buffers and schema changes, so subsequent reset and export can finalize those episodes. diff --git a/source/isaaclab/changelog.d/supermimic.added.rst b/source/isaaclab/changelog.d/supermimic.added.rst new file mode 100644 index 000000000000..3babbcc46a52 --- /dev/null +++ b/source/isaaclab/changelog.d/supermimic.added.rst @@ -0,0 +1 @@ +* Added SuperMimic to the training-library choices reported by the task-listing CLI. diff --git a/source/isaaclab/isaaclab/cli/commands/list_envs.py b/source/isaaclab/isaaclab/cli/commands/list_envs.py index fdba9c8edf1a..d2ea258a99d0 100644 --- a/source/isaaclab/isaaclab/cli/commands/list_envs.py +++ b/source/isaaclab/isaaclab/cli/commands/list_envs.py @@ -73,7 +73,7 @@ def command_list_envs(args: list[str] | None = None) -> None: for index, spec in enumerate(task_specs, start=1): libraries = [ library - for library in ("rl_games", "rlinf", "rsl_rl", "sb3", "skrl", "torchrl") + for library in ("rl_games", "rlinf", "rsl_rl", "sb3", "skrl", "supermimic", "torchrl") if spec.kwargs.get(f"{library}_cfg_entry_point") is not None ] row = [index, spec.id, spec.entry_point, spec.kwargs["env_cfg_entry_point"], ", ".join(libraries) or "(none)"] diff --git a/source/isaaclab/isaaclab/envs/manager_based_rl_env.py b/source/isaaclab/isaaclab/envs/manager_based_rl_env.py index 859cb837cd5b..a7bbe2de7e1e 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_rl_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_rl_env.py @@ -236,10 +236,7 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: self.reset_time_outs = self.termination_manager.time_outs self.reward_buf = self.reward_manager.compute(dt=self.step_dt) - if len(self.recorder_manager.active_terms) > 0: - # update observations for recording if needed - self.obs_buf = self.observation_manager.compute() - self.recorder_manager.record_post_step() + self.recorder_manager.record_post_step() # -- reset envs that terminated/timed-out and log the episode information reset_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1).int() diff --git a/source/isaaclab/isaaclab/envs/mdp/recorders/recorders.py b/source/isaaclab/isaaclab/envs/mdp/recorders/recorders.py index a9315bbca63d..55050f5f43a5 100644 --- a/source/isaaclab/isaaclab/envs/mdp/recorders/recorders.py +++ b/source/isaaclab/isaaclab/envs/mdp/recorders/recorders.py @@ -14,7 +14,7 @@ class InitialStateRecorder(RecorderTerm): """Recorder term that records the initial state of the environment after reset.""" - def record_post_reset(self, env_ids: Sequence[int] | None): + def record_post_reset(self, env_ids: Sequence[int] | slice | None): def extract_env_ids_values(value): nonlocal env_ids if isinstance(value, dict): diff --git a/source/isaaclab/isaaclab/managers/recorder_manager.py b/source/isaaclab/isaaclab/managers/recorder_manager.py index f4894840afb9..58b352b76eaa 100644 --- a/source/isaaclab/isaaclab/managers/recorder_manager.py +++ b/source/isaaclab/isaaclab/managers/recorder_manager.py @@ -54,6 +54,12 @@ class RecorderManagerBaseCfg: export_in_close: bool = False """Whether to export episodes in the close call.""" + buffer_initial_capacity: int = 64 + """Initial samples per world and field. Must be positive; storage grows as needed. + + Set to the expected episode length, including decimation callbacks, to avoid reallocations. + """ + dataset_compression: bool = True """Enable dataset compression.""" @@ -68,24 +74,15 @@ class RecorderTerm(ManagerTermBase): * Post-reset recording: This callback is invoked at the end of `env.reset()`. * Pre-step recording: This callback is invoked at the beginning of `env.step()`, after the step action is processed and before the action is applied by the action manager. - * Post-step recording: This callback is invoked at the end of `env.step()` when all the managers are processed. + * Post-step recording: This callback captures the completed physical transition. In RL environments, + it runs after rewards and terminations, before automatic reset and the policy observation update. """ - def __init__(self, cfg: RecorderTermCfg, env: ManagerBasedEnv): - """Initialize the recorder term. - - Args: - cfg: The configuration object. - env: The environment instance. - """ - # call the base class constructor - super().__init__(cfg, env) - """ User-defined callbacks. """ - def record_pre_reset(self, env_ids: Sequence[int] | None) -> tuple[str | None, torch.Tensor | dict | None]: + def record_pre_reset(self, env_ids: Sequence[int] | slice | None) -> tuple[str | None, torch.Tensor | dict | None]: """Record data at the beginning of env.reset() before reset is effective. Args: @@ -100,7 +97,7 @@ def record_pre_reset(self, env_ids: Sequence[int] | None) -> tuple[str | None, t """ return None, None - def record_post_reset(self, env_ids: Sequence[int] | None) -> tuple[str | None, torch.Tensor | dict | None]: + def record_post_reset(self, env_ids: Sequence[int] | slice | None) -> tuple[str | None, torch.Tensor | dict | None]: """Record data at the end of env.reset(). Args: @@ -122,7 +119,11 @@ def record_pre_step(self) -> tuple[str | None, torch.Tensor | dict | None]: return None, None def record_post_step(self) -> tuple[str | None, torch.Tensor | dict | None]: - """Record data at the end of env.step() when all the managers are processed. + """Record the completed physical transition. + + In RL environments, this runs before automatic reset and observation updates; + ``env.obs_buf`` still holds the preceding policy input. See :ref:`recorder-capture` + for recording current observations. Non-RL environments update observations first. Returns: A tuple of key and value to be recorded. @@ -139,17 +140,6 @@ def record_post_physics_decimation_step(self) -> tuple[str | None, torch.Tensor """ return None, None - def close(self, file_path: str): - """Finalize and "clean up" the recorder term. - - This can include tasks such as appending metadata (e.g. labels) to a file - and properly closing any associated file handles or resources. - - Args: - file_path: the absolute path to the file - """ - pass - class RecorderManager(ManagerBase): """Manager for recording data from recorder terms.""" @@ -163,6 +153,8 @@ def __init__(self, cfg: object, env: ManagerBasedEnv): """ self._term_names: list[str] = [] self._terms: dict[str, RecorderTerm] = {} + self._dataset_file_handler = None + self._failed_episode_dataset_file_handler = None # Do nothing if cfg is None or an empty dict if not cfg: @@ -177,30 +169,28 @@ def __init__(self, cfg: object, env: ManagerBasedEnv): if not isinstance(cfg, RecorderManagerBaseCfg): raise TypeError("Configuration for the recorder manager is not of type RecorderManagerBaseCfg.") - # create episode data buffer indexed by environment id - self._episodes: dict[int, EpisodeData] = dict() - for env_id in range(env.num_envs): - self._episodes[env_id] = EpisodeData() + if cfg.buffer_initial_capacity < 1: + raise ValueError("buffer_initial_capacity must be positive.") + self._fields: list[_RecorderField] = [] + self._env_ids = torch.arange(env.num_envs, device=env.device) + # -1 denotes an episode whose success has not been set. + self._success = torch.full((env.num_envs,), -1, dtype=torch.int8, device=env.device) + self._exported_counts = torch.zeros(2, dtype=torch.long, device=env.device) env_name = getattr(env.cfg, "env_name", None) - self._dataset_file_handler = None if cfg.dataset_export_mode != DatasetExportMode.EXPORT_NONE: self._dataset_file_handler = cfg.dataset_file_handler_class_type() self._dataset_file_handler.create( os.path.join(cfg.dataset_export_dir_path, cfg.dataset_filename), env_name=env_name ) - self._failed_episode_dataset_file_handler = None if cfg.dataset_export_mode == DatasetExportMode.EXPORT_SUCCEEDED_FAILED_IN_SEPARATE_FILES: self._failed_episode_dataset_file_handler = cfg.dataset_file_handler_class_type() self._failed_episode_dataset_file_handler.create( os.path.join(cfg.dataset_export_dir_path, f"{cfg.dataset_filename}_failed"), env_name=env_name ) - self._exported_successful_episode_count = {} - self._exported_failed_episode_count = {} - def __str__(self) -> str: """Returns: A string representation for recorder manager.""" msg = f" contains {len(self._term_names)} active terms.\n" @@ -232,36 +222,14 @@ def active_terms(self) -> list[str]: return self._term_names @property - def exported_successful_episode_count(self, env_id=None) -> int: - """Number of successful episodes. - - Args: - env_id: The environment id. Defaults to None, in which case all environments are considered. - - Returns: - The number of successful episodes. - """ - if not hasattr(self, "_exported_successful_episode_count"): - return 0 - if env_id is not None: - return self._exported_successful_episode_count.get(env_id, 0) - return sum(self._exported_successful_episode_count.values()) + def exported_successful_episode_count(self) -> int: + """Number of successful episodes.""" + return int(self._exported_counts[1].item()) if self.active_terms else 0 @property - def exported_failed_episode_count(self, env_id=None) -> int: - """Number of failed episodes. - - Args: - env_id: The environment id. Defaults to None, in which case all environments are considered. - - Returns: - The number of failed episodes. - """ - if not hasattr(self, "_exported_failed_episode_count"): - return 0 - if env_id is not None: - return self._exported_failed_episode_count.get(env_id, 0) - return sum(self._exported_failed_episode_count.values()) + def exported_failed_episode_count(self) -> int: + """Number of failed episodes.""" + return int(self._exported_counts[0].item()) if self.active_terms else 0 """ Operations. @@ -280,123 +248,94 @@ def reset(self, env_ids: Sequence[int] | slice | None = None) -> dict[str, torch # Do nothing if no active recorder terms are provided if len(self.active_terms) == 0: return {} - # resolve environment ids - if env_ids is None: - env_ids = slice(None) - if isinstance(env_ids, slice): - env_ids = list(range(self._env.num_envs)[env_ids]) - if isinstance(env_ids, torch.Tensor): - env_ids = env_ids.tolist() - for term in self._terms.values(): term.reset(env_ids=env_ids) - - for env_id in env_ids: - self._episodes[env_id] = EpisodeData() - - # nothing to log here + self._reset_buffers(env_ids) return {} - def get_episode(self, env_id: int) -> EpisodeData: - """Returns the episode data for the given environment id. + def get_term(self, name: str) -> RecorderTerm: + """Return the configured recorder term with the given name.""" + return self._terms[name] - Args: - env_id: The environment id. + def get_episode(self, env_id: int) -> EpisodeData: + """Return an owned snapshot of a world's current episode. - Returns: - The episode data for the given environment id. + Tensor leaves have a leading sample dimension. Reading copies data and may synchronize + the device. For live writes, use :meth:`add_to_episodes` and :meth:`set_success_to_episodes`. """ - return self._episodes.get(env_id, EpisodeData()) + episode = EpisodeData() + if not self.active_terms or not 0 <= env_id < self._env.num_envs: + return episode + for field in self._fields: + length = int(field.lengths[env_id].item()) + if length: + data = episode.data + keys = field.name.split("/") + for key in keys[:-1]: + data = data.setdefault(key, {}) + data[keys[-1]] = field.data[:length, env_id].clone() + success = int(self._success[env_id].item()) + episode.success = None if success < 0 else bool(success) + return episode def add_to_episodes(self, key: str, value: torch.Tensor | dict, env_ids: Sequence[int] | slice | None = None): - """Adds the given key-value pair to the episodes for the given environment ids. + """Append a batch to the selected episodes. Args: - key: The key of the given value to be added to the episodes. The key can contain nested keys - separated by '/'. For example, "obs/joint_pos" would add the given value under ['obs']['policy'] - in the underlying dictionary in the episode data. - value: The value to be added to the episodes. The value can be a tensor or a nested dictionary of tensors. - The shape of a tensor in the value is (env_ids, ...). - env_ids: The environment ids. Defaults to None, in which case all environments are considered. + key: Field path, with nested keys separated by '/'. + value: Tensor of shape (selected worlds, ...) or nested dictionary of tensors. + Warp arrays are also accepted. Storage copies the values and preserves dtype/device. + A field's shape, dtype, and device must remain constant while it contains samples. + env_ids: Unique world indices in the same order as the value batch. None selects all worlds. """ - # Do nothing if no active recorder terms are provided - if len(self.active_terms) == 0: - return - - # resolve environment ids - if key is None: + if not self.active_terms or key is None: return - if env_ids is None: - env_ids = slice(None) - if isinstance(env_ids, slice): - env_ids = list(range(self._env.num_envs)[env_ids]) - if isinstance(env_ids, torch.Tensor): - env_ids = env_ids.tolist() - if isinstance(value, dict): for sub_key, sub_value in value.items(): - if isinstance(sub_value, wp.array): - sub_value = wp.to_torch(sub_value) self.add_to_episodes(f"{key}/{sub_key}", sub_value, env_ids) return - if isinstance(value, wp.array): value = wp.to_torch(value) - value = value.clone() # Clone once for all envs - for value_index, env_id in enumerate(env_ids): - if env_id not in self._episodes: - self._episodes[env_id] = EpisodeData() - self._episodes[env_id].env_id = env_id - self._episodes[env_id].add(key, value[value_index], clone=False) + num_selected = ( + self._env.num_envs + if env_ids is None + else (len(range(self._env.num_envs)[env_ids]) if isinstance(env_ids, slice) else len(env_ids)) + ) + if value.ndim < 1 or value.shape[0] != num_selected: + raise ValueError(f"Field '{key}' requires {num_selected} rows, got {tuple(value.shape)}.") + if not num_selected: + return + field = next((field for field in self._fields if field.name == key), None) + if field is None: + # Reject ambiguous leaf/subtree paths before allocating any storage. + if any(key.startswith(f"{f.name}/") or f.name.startswith(f"{key}/") for f in self._fields): + raise ValueError(f"Recorder field '{key}' conflicts with an existing field path.") + field = _RecorderField(key, value, self._env.num_envs, self.cfg.buffer_initial_capacity) + self._fields.append(field) + field.append(value, env_ids) def set_success_to_episodes(self, env_ids: Sequence[int] | slice | None, success_values: torch.Tensor): - """Sets the task success values to the episodes for the given environment ids. - - Args: - env_ids: The environment ids. Defaults to None, in which case all environments are considered. - success_values: The task success values to be set to the episodes. The shape of the tensor is (env_ids, 1). - """ - # Do nothing if no active recorder terms are provided - if len(self.active_terms) == 0: + """Set selected worlds' success flags from a tensor of shape (selected worlds,) or (selected worlds, 1).""" + if not self.active_terms: return - - # resolve environment ids if env_ids is None: env_ids = slice(None) - if isinstance(env_ids, slice): - env_ids = list(range(self._env.num_envs)[env_ids]) - if isinstance(env_ids, torch.Tensor): - env_ids = env_ids.tolist() - - for value_index, env_id in enumerate(env_ids): - self._episodes[env_id].success = success_values[value_index].item() + self._success[env_ids] = success_values.reshape(-1).to(device=self._success.device, dtype=torch.int8) def record_pre_step(self) -> None: """Trigger recorder terms for pre-step functions.""" - # Do nothing if no active recorder terms are provided - if len(self.active_terms) == 0: - return - for term in self._terms.values(): key, value = term.record_pre_step() self.add_to_episodes(key, value) def record_post_step(self) -> None: """Trigger recorder terms for post-step functions.""" - # Do nothing if no active recorder terms are provided - if len(self.active_terms) == 0: - return - for term in self._terms.values(): key, value = term.record_post_step() self.add_to_episodes(key, value) def record_post_physics_decimation_step(self) -> None: """Trigger recorder terms for post-physics step functions in the decimation loop.""" - # Do nothing if no active recorder terms are provided - if len(self.active_terms) == 0: - return - for term in self._terms.values(): key, value = term.record_post_physics_decimation_step() self.add_to_episodes(key, value) @@ -413,17 +352,13 @@ def record_pre_reset(self, env_ids: Sequence[int] | slice | None, force_export_o if env_ids is None: env_ids = slice(None) - if isinstance(env_ids, slice): - env_ids = list(range(self._env.num_envs)[env_ids]) - if isinstance(env_ids, torch.Tensor): - env_ids = env_ids.tolist() for term in self._terms.values(): key, value = term.record_pre_reset(env_ids) self.add_to_episodes(key, value, env_ids) # Set task success values for the relevant episodes - success_results = torch.zeros(len(env_ids), dtype=bool, device=self._env.device) + success_results = torch.zeros_like(self._success[env_ids], dtype=torch.bool) # Check success indicator from termination terms if hasattr(self._env, "termination_manager"): if "success" in self._env.termination_manager.active_terms: @@ -443,8 +378,8 @@ def record_post_reset(self, env_ids: Sequence[int] | slice | None) -> None: if len(self.active_terms) == 0: return - if isinstance(env_ids, slice): - env_ids = list(range(self._env.num_envs)[env_ids]) + if env_ids is None: + env_ids = slice(None) for term in self._terms.values(): key, value = term.record_post_reset(env_ids) @@ -485,11 +420,11 @@ def export_episodes( return if env_ids is None: - env_ids = slice(None) - if isinstance(env_ids, slice): - env_ids = list(range(self._env.num_envs)[env_ids]) - if isinstance(env_ids, torch.Tensor): - env_ids = env_ids.tolist() + env_ids = self._env_ids + elif isinstance(env_ids, slice): + env_ids = self._env_ids[env_ids] + else: + env_ids = torch.as_tensor(env_ids, dtype=torch.long, device=self._env_ids.device) # Handle demo_ids processing if demo_ids is not None: @@ -502,58 +437,46 @@ def export_episodes( duplicates = [x for i, x in enumerate(demo_ids) if demo_ids.index(x) != i] raise ValueError(f"demo_ids must be unique. Found duplicates: {list(set(duplicates))}") - # Export episode data through dataset exporter - need_to_flush = False + present = torch.zeros(self._env.num_envs, dtype=torch.bool, device=self._env_ids.device) + for field in self._fields: + present |= (field.lengths > 0).to(present.device) + succeeded = self._success == 1 - if any(env_id in self._episodes and not self._episodes[env_id].is_empty() for env_id in env_ids): + # Only file export materializes per-world snapshots. EXPORT_NONE stays batched. + need_to_flush = False + mode = self.cfg.dataset_export_mode + if mode != DatasetExportMode.EXPORT_NONE and present[env_ids].any().item(): ep_meta = self.get_ep_meta() - if self._dataset_file_handler is not None: - self._dataset_file_handler.add_env_args(ep_meta) - if self._failed_episode_dataset_file_handler is not None: - self._failed_episode_dataset_file_handler.add_env_args(ep_meta) - - for i, env_id in enumerate(env_ids): - if env_id in self._episodes and not self._episodes[env_id].is_empty(): - self._episodes[env_id].pre_export() - - episode_succeeded = self._episodes[env_id].success - target_dataset_file_handler = None - if (self.cfg.dataset_export_mode == DatasetExportMode.EXPORT_ALL) or ( - self.cfg.dataset_export_mode == DatasetExportMode.EXPORT_SUCCEEDED_ONLY and episode_succeeded - ): - target_dataset_file_handler = self._dataset_file_handler - elif self.cfg.dataset_export_mode == DatasetExportMode.EXPORT_SUCCEEDED_FAILED_IN_SEPARATE_FILES: - if episode_succeeded: - target_dataset_file_handler = self._dataset_file_handler - else: - target_dataset_file_handler = self._failed_episode_dataset_file_handler - if target_dataset_file_handler is not None: - # Use corresponding demo_id if provided, otherwise None - current_demo_id = demo_ids[i] if demo_ids is not None else None - target_dataset_file_handler.write_episode( - self._episodes[env_id], current_demo_id, self.cfg.dataset_compression - ) - need_to_flush = True - # Update episode count - if episode_succeeded: - self._exported_successful_episode_count[env_id] = ( - self._exported_successful_episode_count.get(env_id, 0) + 1 - ) - else: - self._exported_failed_episode_count[env_id] = self._exported_failed_episode_count.get(env_id, 0) + 1 - # Reset the episode buffer for the given environment after export - self._episodes[env_id] = EpisodeData() - + for handler in (self._dataset_file_handler, self._failed_episode_dataset_file_handler): + if handler is not None: + handler.add_env_args(ep_meta) + selected = env_ids.tolist() + selected_present = present[env_ids].tolist() + selected_success = succeeded[env_ids].tolist() + for i, env_id in enumerate(selected): + if not selected_present[i]: + continue + handler = self._dataset_file_handler + if mode == DatasetExportMode.EXPORT_SUCCEEDED_FAILED_IN_SEPARATE_FILES: + if not selected_success[i]: + handler = self._failed_episode_dataset_file_handler + elif mode == DatasetExportMode.EXPORT_SUCCEEDED_ONLY and not selected_success[i]: + continue + episode = self.get_episode(env_id) + demo_id = demo_ids[i] if demo_ids is not None else None + handler.write_episode(episode, demo_id, self.cfg.dataset_compression) + need_to_flush = True + + self._exported_counts[0] += (present & ~succeeded)[env_ids].sum() + self._exported_counts[1] += (present & succeeded)[env_ids].sum() + self._reset_buffers(env_ids) if need_to_flush: - if self._dataset_file_handler is not None: - self._dataset_file_handler.flush() - if self._failed_episode_dataset_file_handler is not None: - self._failed_episode_dataset_file_handler.flush() + for handler in (self._dataset_file_handler, self._failed_episode_dataset_file_handler): + if handler is not None: + handler.flush() def close(self): - """Closes the recorder manager by exporting any remaining data to file as well as properly - closes the recorder terms. - """ + """Export remaining episodes if configured and close dataset files.""" # Do nothing if no active recorder terms are provided if len(self.active_terms) == 0: return @@ -563,13 +486,17 @@ def close(self): self._dataset_file_handler.close() if self._failed_episode_dataset_file_handler is not None: self._failed_episode_dataset_file_handler.close() - for term in self._terms.values(): - term.close(os.path.join(self.cfg.dataset_export_dir_path, self.cfg.dataset_filename)) """ Helper functions. """ + def _reset_buffers(self, env_ids): + """Discard selected samples by resetting indices; retain allocations for the next episodes.""" + for field in self._fields: + field.reset(env_ids) + self._success[slice(None) if env_ids is None else env_ids] = -1 + def _prepare_terms(self): """Prepares a list of recorder terms.""" # check if config is dict already @@ -587,6 +514,7 @@ def _prepare_terms(self): "export_in_record_pre_reset", "export_in_close", "dataset_compression", + "buffer_initial_capacity", ]: continue # check if term config is None @@ -605,3 +533,65 @@ def _prepare_terms(self): # add term name and parameters self._term_names.append(term_name) self._terms[term_name] = term + + +class _RecorderField: + """One field's time-major allocation and independent per-world insertion indices. + + No per-world Python containers or tensor views are created during append. The host upper + bound avoids reading device indices in steady state; only capacity exhaustion checks synchronize. + """ + + def __init__(self, name: str, value: torch.Tensor, num_envs: int, capacity: int): + self.name = name + # Retained storage outlives the caller's inference context: reset/export may run outside it. + with torch.inference_mode(False): + self.data = torch.empty((capacity, num_envs, *value.shape[1:]), dtype=value.dtype, device=value.device) + self.lengths = torch.zeros(num_envs, dtype=torch.long, device=value.device) + self.worlds = torch.arange(num_envs, device=value.device) + self.upper_bound = 0 + + @torch.no_grad() + def append(self, value: torch.Tensor, env_ids): + if value.shape[1:] != self.data.shape[2:] or value.dtype != self.data.dtype or value.device != self.data.device: + if self.lengths.any().item(): + raise ValueError(f"Recorder field '{self.name}' changed shape, dtype, or device with active samples.") + with torch.inference_mode(False): + self.data = torch.empty( + (*self.data.shape[:2], *value.shape[1:]), dtype=value.dtype, device=value.device + ) + self.lengths = self.lengths.to(value.device) + self.worlds = self.worlds.to(value.device) + self.upper_bound = 0 + if env_ids is None: + worlds, indices = self.worlds, self.lengths + else: + if not isinstance(env_ids, slice): + env_ids = torch.as_tensor(env_ids, dtype=torch.long, device=self.data.device) + worlds = self.worlds[env_ids] + indices = self.lengths[worlds] + capacity = self.data.shape[0] + if self.upper_bound >= capacity: + self.upper_bound = int(self.lengths.max().item()) + if self.upper_bound >= capacity: + with torch.inference_mode(False): + data = self.data.new_empty((2 * capacity, *self.data.shape[1:])) + data[:capacity].copy_(self.data) + self.data = data + self.data[indices, worlds] = value + if env_ids is None: + self.lengths.add_(1) + else: + self.lengths[worlds] = indices + 1 + self.upper_bound += 1 + + def reset(self, env_ids): + if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)): + self.lengths.zero_() + self.upper_bound = 0 + else: + if not isinstance(env_ids, slice): + env_ids = torch.as_tensor(env_ids, dtype=torch.long, device=self.data.device) + self.lengths[env_ids] = 0 + if not isinstance(env_ids, slice) and len(env_ids) == len(self.lengths): + self.upper_bound = 0 diff --git a/source/isaaclab/isaaclab/utils/datasets/episode_data.py b/source/isaaclab/isaaclab/utils/datasets/episode_data.py index bb528d4c630a..c401906b3aa4 100644 --- a/source/isaaclab/isaaclab/utils/datasets/episode_data.py +++ b/source/isaaclab/isaaclab/utils/datasets/episode_data.py @@ -9,7 +9,11 @@ class EpisodeData: - """Class to store episode data.""" + """An episode's tensor data and playback cursors. + + Data leaves are tensors with a leading sample dimension. Collection belongs to the recorder + manager; this container holds snapshots or loaded datasets without accumulating samples. + """ def __init__(self) -> None: """Initializes episode data class.""" @@ -85,34 +89,6 @@ def is_empty(self): """Check if the episode data is empty.""" return not self._data - def add(self, key: str, value: torch.Tensor | dict, clone: bool = True): - """Add a key-value pair to the dataset. - - The key can be nested by using the "/" character. - For example: "obs/joint_pos". - - Args: - key: The key name. - value: The corresponding value of tensor type or of dict type. - clone: Whether to clone the tensor value before storing it in the episode data. - """ - if isinstance(value, dict): - for sub_key, sub_value in value.items(): - self.add(f"{key}/{sub_key}", sub_value, clone=clone) - return - - stored = value.clone() if (clone and isinstance(value, torch.Tensor)) else value - sub_keys = key.split("/") - data = self._data - for sub_key in sub_keys[:-1]: - if sub_key not in data: - data[sub_key] = {} - data = data[sub_key] - # Accumulate in lists to avoid copying tensors on each append. - if sub_keys[-1] not in data: - data[sub_keys[-1]] = [] - data[sub_keys[-1]].append(stored) - def get_initial_state(self) -> torch.Tensor | None: """Get the initial state from the dataset.""" return self._data.get("initial_state") @@ -158,18 +134,6 @@ def get_next_joint_target(self) -> dict | torch.Tensor | None: self._next_joint_target_index += 1 return joint_target - def pre_export(self): - """Prepare data for export by converting lists to tensors.""" - - def pre_export_helper(data): - for key, value in data.items(): - if isinstance(value, list): - data[key] = torch.stack(value) - elif isinstance(value, dict): - pre_export_helper(value) - - pre_export_helper(self._data) - def _index_nested(data: dict | torch.Tensor, index: int, keep_dim: bool) -> dict | torch.Tensor | None: """Select ``index`` along the first dimension of every tensor in a (nested) dict. diff --git a/source/isaaclab/isaaclab/utils/math.py b/source/isaaclab/isaaclab/utils/math.py index c3366a3d8f81..032c815d7807 100644 --- a/source/isaaclab/isaaclab/utils/math.py +++ b/source/isaaclab/isaaclab/utils/math.py @@ -495,7 +495,7 @@ def axis_angle_from_quat(quat: torch.Tensor, eps: float = 1.0e-6) -> torch.Tenso angle = 2.0 * half_angle # check whether to apply Taylor approximation sin_half_angles_over_angles = torch.where( - angle.abs() > eps, torch.sin(half_angle) / angle, 0.5 - angle * angle / 48 + angle.abs() > eps, torch.sin(half_angle) / angle.clamp_min(eps), 0.5 - angle * angle / 48 ) return quat[..., :3] / sin_half_angles_over_angles.unsqueeze(-1) diff --git a/source/isaaclab/test/cli/test_replay_demos_loop_termination.py b/source/isaaclab/test/cli/test_replay_demos_loop_termination.py index 0692374b59ea..542f42469d90 100644 --- a/source/isaaclab/test/cli/test_replay_demos_loop_termination.py +++ b/source/isaaclab/test/cli/test_replay_demos_loop_termination.py @@ -77,7 +77,7 @@ def __init__(self, actions: torch.Tensor): def load_episode(self, episode_name, device): episode = EpisodeData() - episode.data = {"initial_state": {}, "actions": list(self._actions)} + episode.data = {"initial_state": {}, "actions": self._actions} return episode diff --git a/source/isaaclab/test/envs/test_manager_based_rl_env_unit.py b/source/isaaclab/test/envs/test_manager_based_rl_env_unit.py index e8042d488292..0a45428d7daf 100644 --- a/source/isaaclab/test/envs/test_manager_based_rl_env_unit.py +++ b/source/isaaclab/test/envs/test_manager_based_rl_env_unit.py @@ -12,8 +12,20 @@ import gymnasium as gym import numpy as np import pytest +import torch from isaaclab.envs import ManagerBasedRLEnv +from isaaclab.managers import ( + DatasetExportMode, + ObservationGroupCfg, + ObservationManager, + ObservationTermCfg, + RecorderManager, + RecorderManagerBaseCfg, + RecorderTerm, + RecorderTermCfg, +) +from isaaclab.utils.noise import UniformNoiseCfg pytestmark = pytest.mark.unit @@ -102,3 +114,99 @@ def test_obs_space_follows_clip_constraint(): assert term_space.shape == expected_shapes[term_name] assert np.all(term_space.low == low) assert np.all(term_space.high == high) + + +def _state_observation(env): + return env.scene.state + + +class _PhysicalRecorder(RecorderTerm): + def record_post_step(self): + self.latest = self._env.scene.state.clone() + return None, None + + +class _ObservationRecorder(RecorderTerm): + def record_post_step(self): + self.latest = self._env.observation_manager.compute_group("state", update_history=False) + return None, None + + +def _make_stepping_env(recorder_type=None): + """Exercise the real step/observation/recorder owners with a minimal deterministic physics leaf.""" + env = object.__new__(ManagerBasedRLEnv) + env._is_closed = True + env.cfg = SimpleNamespace(decimation=1, sim=SimpleNamespace(dt=0.02, render_interval=1), compute_final_obs=False) + env.scene = SimpleNamespace(num_envs=3, state=torch.zeros(3, 2)) + env.scene.write_data_to_sim = lambda: None + env.scene.update = lambda dt: None + env.sim = SimpleNamespace( + device="cpu", + is_playing=lambda: True, + is_rendering=False, + step=lambda render: env.scene.state.add_(1), + consume_reset_request=lambda: False, + ) + env._physics_handles_decimation = False + env._sim_step_counter = env.common_step_counter = 0 + env.episode_length_buf = torch.zeros(3, dtype=torch.long) + env.render_enabled = False + env.video_recorders = [] + env.extras = {} + env.action_manager = SimpleNamespace(process_action=lambda action: None, apply_action=lambda: None) + mask = torch.zeros(3, dtype=torch.bool) + + def terminate(): + mask[0] = env.common_step_counter == 4 + return mask + + env.termination_manager = SimpleNamespace( + active_terms=[], compute=terminate, terminated=mask, time_outs=torch.zeros_like(mask) + ) + env.reward_manager = SimpleNamespace(compute=lambda dt: torch.zeros(3)) + env.command_manager = SimpleNamespace(compute=lambda dt: None) + env.event_manager = SimpleNamespace(active_terms=[]) + group = ObservationGroupCfg(history_length=3, flatten_history_dim=False, enable_corruption=True) + group.state = ObservationTermCfg(func=_state_observation, noise=UniformNoiseCfg(n_min=-0.1, n_max=0.1)) + state = ObservationGroupCfg(enable_corruption=False) + state.state = ObservationTermCfg(func=_state_observation) + env.observation_manager = ObservationManager({"policy": group, "state": state}, env) + cfg = RecorderManagerBaseCfg(dataset_export_mode=DatasetExportMode.EXPORT_NONE) + if recorder_type is not None: + cfg.capture = RecorderTermCfg(class_type=recorder_type) + env.recorder_manager = RecorderManager(cfg, env) + + def reset(env_ids): + env.scene.state[env_ids] = -100.0 + env.observation_manager.reset(env_ids) + env.recorder_manager.reset(env_ids) + + env._reset_idx = reset + env.obs_buf = env.observation_manager.compute(update_history=True) + return env + + +def test_raw_capture_preserves_observation_noise_history_and_rng(): + """Capture sees terminal state without implicit observation work or changes to policy history/RNG.""" + outputs, states = [], [] + for recorder_type in (None, _PhysicalRecorder, _ObservationRecorder): + torch.manual_seed(17) + env = _make_stepping_env(recorder_type) + history = [env.step(torch.empty(3, 0))[0]["policy"].clone() for _ in range(4)] + outputs.append(torch.stack(history)) + states.append(torch.get_rng_state()) + if recorder_type is not None: + term = env.recorder_manager.get_term("capture") + assert not env.recorder_manager._fields + # Both direct capture and explicit noiseless observations see the state before reset. + assert torch.equal(term.latest, torch.full((3, 2), 4.0)) + assert torch.equal(env.scene.state[0], torch.full((2,), -100.0)) + env.recorder_manager.close() + assert torch.equal(outputs[0], outputs[1]) + assert torch.equal(states[0], states[1]) + assert torch.equal(outputs[0], outputs[2]) + assert torch.equal(states[0], states[2]) + # Continuing worlds retain steps 2, 3 and 4; the reset world's history starts afresh. + expected = torch.tensor([2.0, 3.0, 4.0])[None, :, None] + assert torch.all((outputs[1][-1, 1:] - expected).abs() <= 0.1) + assert torch.all((outputs[1][-1, 0] + 100.0).abs() <= 0.1) diff --git a/source/isaaclab/test/managers/test_recorder_manager.py b/source/isaaclab/test/managers/test_recorder_manager.py index c07dfba99d10..e0396da5543e 100644 --- a/source/isaaclab/test/managers/test_recorder_manager.py +++ b/source/isaaclab/test/managers/test_recorder_manager.py @@ -4,18 +4,21 @@ # SPDX-License-Identifier: BSD-3-Clause from __future__ import annotations -import os -import shutil -import tempfile -import uuid +import ast +import inspect +import sys from collections.abc import Sequence +from types import SimpleNamespace from typing import TYPE_CHECKING, cast +import h5py import pytest import torch +import warp as wp from isaaclab.managers import DatasetExportMode, RecorderManager, RecorderManagerBaseCfg, RecorderTerm, RecorderTermCfg from isaaclab.utils import configclass +from isaaclab.utils.datasets import EpisodeData pytestmark = pytest.mark.unit @@ -26,11 +29,13 @@ class DummyResetRecorderTerm(RecorderTerm): """Dummy recorder term that records dummy data.""" - def record_pre_reset(self, env_ids: Sequence[int] | None) -> tuple[str | None, torch.Tensor | None]: - return "record_pre_reset", torch.ones(self._env.num_envs, 2, device=self._env.device) + def record_pre_reset(self, env_ids: Sequence[int] | slice | None) -> tuple[str, torch.Tensor]: + self.env_ids = env_ids + return "record_pre_reset", torch.ones(self._env.num_envs, 2, device=self._env.device)[env_ids] - def record_post_reset(self, env_ids: Sequence[int] | None) -> tuple[str | None, torch.Tensor | None]: - return "record_post_reset", torch.ones(self._env.num_envs, 3, device=self._env.device) + def record_post_reset(self, env_ids: Sequence[int] | slice | None) -> tuple[str, torch.Tensor]: + self.env_ids = env_ids + return "record_post_reset", torch.ones(self._env.num_envs, 3, device=self._env.device)[env_ids] class DummyStepRecorderTerm(RecorderTerm): @@ -47,137 +52,195 @@ def record_post_step(self) -> tuple[str | None, torch.Tensor | None]: class DummyRecorderManagerCfg(RecorderManagerBaseCfg): """Dummy recorder configurations.""" - @configclass - class DummyResetRecorderTermCfg(RecorderTermCfg): - """Configuration for the dummy reset recorder term.""" - - class_type: type[RecorderTerm] = DummyResetRecorderTerm - - @configclass - class DummyStepRecorderTermCfg(RecorderTermCfg): - """Configuration for the dummy step recorder term.""" - - class_type: type[RecorderTerm] = DummyStepRecorderTerm - - record_reset_term = DummyResetRecorderTermCfg() - record_step_term = DummyStepRecorderTermCfg() + record_reset_term = RecorderTermCfg(class_type=DummyResetRecorderTerm) + record_step_term = RecorderTermCfg(class_type=DummyStepRecorderTerm) dataset_export_mode = DatasetExportMode.EXPORT_ALL -@configclass -class DummyEnvCfg: - """Dummy environment configuration.""" - - @configclass - class DummySimCfg: - """Configuration for the dummy sim.""" - - dt: float = 0.01 - render_interval: int = 1 - - @configclass - class DummySceneCfg: - """Configuration for the dummy scene.""" - - num_envs: int = 1 - - decimation: int = 1 - sim: DummySimCfg = DummySimCfg() - scene: DummySceneCfg = DummySceneCfg() - - -class DummySimulation: - """Minimal simulation double used by :class:`RecorderManager`.""" - - def is_playing(self) -> bool: - """Return whether the simulated timeline is playing.""" - return True - - -class DummyTerminationManager: - """Minimal termination manager double without success terms.""" - - active_terms: list[str] = [] - - -class DummyEnv: - """Minimal environment double used by recorder terms and metadata export.""" - - def __init__(self, device: str = "cpu", num_envs: int = 20) -> None: - self.num_envs = num_envs - self.device = device - self.sim = DummySimulation() - self.cfg = DummyEnvCfg() - self.cfg.scene.num_envs = num_envs - self.termination_manager = DummyTerminationManager() - - -def create_dummy_env(device: str = "cpu", num_envs: int = 20) -> ManagerBasedEnv: +def create_dummy_env(device: str = "cpu", num_envs: int = 3) -> ManagerBasedEnv: """Create a minimal environment double.""" - - return cast("ManagerBasedEnv", DummyEnv(device=device, num_envs=num_envs)) - - -@pytest.fixture -def dataset_dir(): - """Create directory to dump results.""" - test_dir = tempfile.mkdtemp() - yield test_dir - # Cleanup - shutil.rmtree(test_dir) - - -def test_str(dataset_dir): - """Test the string representation of the recorder manager.""" - # create recorder manager - cfg = DummyRecorderManagerCfg() - cfg.dataset_export_dir_path = dataset_dir - cfg.dataset_filename = f"{uuid.uuid4()}.hdf5" - recorder_manager = RecorderManager(cfg, create_dummy_env()) - assert len(recorder_manager.active_terms) == 2 - manager_str = str(recorder_manager) - assert "contains 2 active terms" in manager_str - assert "record_reset_term" in manager_str - assert "record_step_term" in manager_str - # the dataset file is created on construction - assert os.path.exists(os.path.join(cfg.dataset_export_dir_path, cfg.dataset_filename)) - recorder_manager.close() - - -def test_record(dataset_dir): - """Test the recording of the data.""" + env = SimpleNamespace( + num_envs=num_envs, + device=device, + sim=SimpleNamespace(is_playing=lambda: True), + cfg=SimpleNamespace( + decimation=1, + sim=SimpleNamespace(dt=0.01, render_interval=1), + scene=SimpleNamespace(num_envs=num_envs), + ), + termination_manager=SimpleNamespace(active_terms=[]), + ) + return cast("ManagerBasedEnv", env) + + +def test_record(tmp_path): + """Capture callbacks and full/partial resets preserve episode boundaries and selectors.""" env = create_dummy_env() - # create recorder manager - cfg = DummyRecorderManagerCfg() - cfg.dataset_export_dir_path = dataset_dir - cfg.dataset_filename = f"{uuid.uuid4()}.hdf5" - recorder_manager = RecorderManager(cfg, env) - - # record the step data - recorder_manager.record_pre_step() - recorder_manager.record_post_step() - - recorder_manager.record_pre_step() - recorder_manager.record_post_step() - - # check the recorded data - for env_id in range(env.num_envs): - episode = recorder_manager.get_episode(env_id) - assert torch.stack(episode.data["record_pre_step"]).shape == (2, 4) - assert torch.stack(episode.data["record_post_step"]).shape == (2, 5) - - # Trigger pre-reset callbacks which then export and clean the episode data - recorder_manager.record_pre_reset(env_ids=slice(None)) - for env_id in range(env.num_envs): - episode = recorder_manager.get_episode(env_id) - assert episode.is_empty() - - recorder_manager.record_post_reset(env_ids=slice(None)) + for cfg in (None, {}, RecorderManagerBaseCfg()): + disabled = RecorderManager(cfg, env) + assert disabled.exported_successful_episode_count == disabled.exported_failed_episode_count == 0 + cfg = DummyRecorderManagerCfg(dataset_export_dir_path=str(tmp_path), dataset_filename="dataset.hdf5") + manager = RecorderManager(cfg, env) + assert (tmp_path / cfg.dataset_filename).is_file() + reset_term = manager.get_term("record_reset_term") + + for _ in range(2): + manager.record_pre_step() + manager.record_post_step() for env_id in range(env.num_envs): - episode = recorder_manager.get_episode(env_id) - assert torch.stack(episode.data["record_post_reset"]).shape == (1, 3) - recorder_manager.reset(slice(1, None, 2)) + episode = manager.get_episode(env_id) + assert episode.data["record_pre_step"].shape == (2, 4) + assert episode.data["record_post_step"].shape == (2, 5) + + manager.record_pre_reset(slice(None)) # Capture, export and clear the completed episodes. + assert reset_term.env_ids == slice(None) + assert all(manager.get_episode(i).is_empty() for i in range(env.num_envs)) + manager.record_post_reset(None) + assert reset_term.env_ids == slice(None) + manager.reset(slice(1, None, 2)) for env_id in range(env.num_envs): - assert recorder_manager.get_episode(env_id).is_empty() == (env_id % 2 == 1) - recorder_manager.close() + assert manager.get_episode(env_id).is_empty() == (env_id % 2 == 1) + manager.record_post_reset(slice(1, None, 2)) + assert reset_term.env_ids == slice(1, None, 2) + assert all(manager.get_episode(i).data["record_post_reset"].shape == (1, 3) for i in range(env.num_envs)) + manager.close() + + +def test_batched_episode_indices_and_ownership(monkeypatch): + """Asynchronous resets, field cadence, growth and snapshots must preserve exact samples.""" + cfg = DummyRecorderManagerCfg(dataset_export_mode=DatasetExportMode.EXPORT_NONE, buffer_initial_capacity=2) + manager = RecorderManager(cfg, create_dummy_env(num_envs=3)) + manager.add_to_episodes("unused", torch.empty(0, 7), []) + assert not manager._fields + with pytest.raises(ValueError, match="requires 3 rows"): + manager.add_to_episodes("invalid", torch.empty(2, 7)) + assert not manager._fields + source = torch.tensor([[1.0, -0.0], [2.0, float("nan")], [3.0, 4.0]]) + # Policies commonly collect under inference_mode, then reset/export after leaving it. + with torch.inference_mode(): + manager.add_to_episodes("obs", {"state": source}) + wp.init() + manager.add_to_episodes("initial_state", wp.from_torch(torch.tensor([[10], [20], [30]], dtype=torch.int32))) + source.fill_(99) + manager.add_to_episodes("obs/state", source) + snapshot = manager.get_episode(1) + assert snapshot.data["obs"]["state"][0, 0] == 2 + assert torch.isnan(snapshot.data["obs"]["state"][0, 1]) + assert torch.signbit(manager.get_episode(0).data["obs"]["state"][0, 1]) + assert snapshot.data["initial_state"].dtype == torch.int32 + assert snapshot.success is None + + manager.reset(torch.tensor([1])) + with torch.inference_mode(): + manager.add_to_episodes("obs/state", torch.tensor([[42.0, 43.0], [52.0, 53.0]]), [2, 1]) + manager.add_to_episodes("obs/state", torch.tensor([[62.0, 63.0]]), slice(0, 1)) + assert torch.equal(manager.get_episode(1).data["obs"]["state"], torch.tensor([[52.0, 53.0]])) + assert manager.get_episode(1).data.keys() == {"obs"} + assert torch.equal(manager.get_episode(2).data["obs"]["state"][:, 0], torch.tensor([3.0, 99.0, 42.0])) + assert torch.equal(manager.get_episode(0).data["obs"]["state"][:, 0], torch.tensor([1.0, 99.0, 62.0])) + assert snapshot.data["obs"]["state"].shape == (2, 2) + snapshot = manager.get_episode(1) + snapshot.data["obs"]["state"].fill_(-100) + assert manager.get_episode(1).data["obs"]["state"][0, 0] == 52 + + with pytest.raises(ValueError, match="changed shape"): + manager.add_to_episodes("obs/state", torch.zeros(3, 5)) + with pytest.raises(ValueError, match="conflicts"): + manager.add_to_episodes("obs", torch.zeros(3, 2)) + manager.add_to_episodes("obs/state", torch.empty(0, 2), []) + assert manager.get_episode(0).data["obs"]["state"].shape == (3, 2) + + manager.reset() + # A new schema is allowed once every world has released this field, e.g. manual annotation. + with torch.inference_mode(): + manager.add_to_episodes("obs/state", torch.zeros(3, 5, dtype=torch.bool)) + assert manager.get_episode(0).data["obs"]["state"].dtype == torch.bool + manager.reset() + + def forbidden(*args, **kwargs): + raise AssertionError("Steady-state recording must reuse storage without episode objects or host index reads") + + # Architectural/performance gate: preallocated insertion and reset need neither episode objects, + # per-world container appends, nor device-to-host scalar reads. Export-none also stays batched. + with monkeypatch.context() as patch: + patch.setattr(EpisodeData, "__init__", forbidden) + patch.setattr(torch.Tensor, "item", forbidden) + patch.setattr(torch, "empty", forbidden) + manager.add_to_episodes("obs/state", torch.zeros(3, 5, dtype=torch.bool)) + manager.set_success_to_episodes([0, 2], torch.tensor([[True], [False]])) + manager.export_episodes([0, 2]) + manager.reset([1]) + assert manager.exported_successful_episode_count == 1 + assert manager.exported_failed_episode_count == 1 + assert all(manager.get_episode(i).is_empty() for i in range(3)) + manager.close() + + +@pytest.mark.parametrize("mode", list(DatasetExportMode)) +def test_export_batched_episodes(tmp_path, mode): + """Export modes, custom IDs, nested tensor schema and counters survive partial finalization.""" + cfg = DummyRecorderManagerCfg(dataset_export_dir_path=str(tmp_path), dataset_export_mode=mode, export_in_close=True) + manager = RecorderManager(cfg, create_dummy_env(num_envs=3)) + values = torch.tensor([[1.0, -0.0], [2.0, float("nan")], [3.0, 4.0]]) + manager.add_to_episodes("actions", values) + manager.add_to_episodes("obs", {"flag": torch.tensor([True, False, True])}) + manager.set_success_to_episodes(None, torch.tensor([[True], [False], [True]])) + with pytest.raises(ValueError, match="must match"): + manager.export_episodes([2, 1], [30]) + with pytest.raises(ValueError, match="must be unique"): + manager.export_episodes([2, 1], [30, 30]) + manager.export_episodes(torch.tensor([2, 1]), [30, 10]) + assert manager.get_episode(1).is_empty() + assert manager.get_episode(2).is_empty() + assert manager.get_episode(0).success is True + manager.export_episodes(slice(0, 1), [20]) + manager.export_episodes() # Empty episodes do not affect counters. + assert manager.exported_successful_episode_count == 2 + assert manager.exported_failed_episode_count == 1 + manager.close() + if mode == DatasetExportMode.EXPORT_NONE: + assert list(tmp_path.iterdir()) == [] + return + with h5py.File(tmp_path / "dataset.hdf5") as file: + expected = {"demo_30", "demo_20"} + if mode == DatasetExportMode.EXPORT_ALL: + expected.add("demo_10") + assert set(file["data"]) == expected + assert file["data/demo_20/actions"][()].tobytes() == values[0:1].numpy().tobytes() + assert bool(file["data/demo_30/obs/flag"][0]) + assert bool(file["data/demo_20"].attrs["success"]) + if mode == DatasetExportMode.EXPORT_SUCCEEDED_FAILED_IN_SEPARATE_FILES: + with h5py.File(tmp_path / "dataset_failed.hdf5") as file: + assert set(file["data"]) == {"demo_10"} + assert file["data/demo_10/actions"][()].tobytes() == values[1:2].numpy().tobytes() + assert not bool(file["data/demo_10"].attrs["success"]) + + +def test_recorder_storage_boundary(): + """Keep storage independent of consumers, with batched ownership and no Python world loops.""" + import isaaclab.managers.recorder_manager as module + + source = ast.parse(inspect.getsource(module)) + # Batched buffers own collection; snapshots and capture terms have no second writer lifecycle. + assert not {"add", "pre_export"} & vars(EpisodeData).keys() + assert "__init__" not in vars(module.RecorderTerm) + assert not hasattr(module.RecorderTerm, "close") + # Observation evaluation belongs to explicit terms, never generic capture dispatch. + assert not hasattr(module.RecorderTerm, "requires_post_step_observations") + assert not hasattr(module.RecorderManager, "requires_post_step_observations") + assert not any( + isinstance(node, ast.Attribute) and node.attr in {"observation_manager", "_episodes"} + for node in ast.walk(source) + ) + allowed = sys.stdlib_module_names | {"isaaclab", "torch", "warp", "prettytable"} + for node in ast.walk(source): + if isinstance(node, ast.Import): + assert all(alias.name.split(".")[0] in allowed for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.level == 0: + assert (node.module or "").split(".")[0] in allowed + elif isinstance(node, ast.FunctionDef) and node.name.startswith("exported_"): + assert not any(isinstance(child, ast.Name) and child.id == "hasattr" for child in ast.walk(node)) + field = next(node for node in source.body if isinstance(node, ast.ClassDef) and node.name == "_RecorderField") + append = next(node for node in field.body if isinstance(node, ast.FunctionDef) and node.name == "append") + assert not any(isinstance(node, (ast.For, ast.While, ast.comprehension)) for node in ast.walk(append)) diff --git a/source/isaaclab/test/utils/test_episode_data.py b/source/isaaclab/test/utils/test_episode_data.py deleted file mode 100644 index 1d9536dc05ac..000000000000 --- a/source/isaaclab/test/utils/test_episode_data.py +++ /dev/null @@ -1,147 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause -import pytest -import torch - -from isaaclab.test.utils import DeviceScope, test_devices -from isaaclab.utils.datasets import EpisodeData - -pytestmark = pytest.mark.unit - - -@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) -def test_is_empty(device): - """Test checking whether the episode is empty.""" - episode = EpisodeData() - assert episode.is_empty() - - episode.add("key", torch.tensor([1, 2, 3], device=device)) - assert not episode.is_empty() - - -@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) -def test_add_tensors(device): - """Test appending tensor data to the episode.""" - dummy_data_0 = torch.tensor([0], device=device) - dummy_data_1 = torch.tensor([1], device=device) - expected_added_data = torch.cat((dummy_data_0.unsqueeze(0), dummy_data_1.unsqueeze(0))) - episode = EpisodeData() - - # test adding data to a key that does not exist - episode.add("key", dummy_data_0) - key_data = torch.stack(episode.data.get("key")) - assert key_data is not None - assert torch.equal(key_data, dummy_data_0.unsqueeze(0)) - - # test adding data to a key that exists - episode.add("key", dummy_data_1) - key_data = torch.stack(episode.data.get("key")) - assert key_data is not None - assert torch.equal(key_data, expected_added_data) - - # test adding data to a key with "/" in the name - episode.add("first/second", dummy_data_0) - first_data = episode.data.get("first") - assert first_data is not None - second_data = torch.stack(first_data.get("second")) - assert second_data is not None - assert torch.equal(second_data, dummy_data_0.unsqueeze(0)) - - # test adding data to a key with "/" in the name that already exists - episode.add("first/second", dummy_data_1) - first_data = episode.data.get("first") - assert first_data is not None - second_data = torch.stack(first_data.get("second")) - assert second_data is not None - assert torch.equal(second_data, expected_added_data) - - -@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) -def test_add_dict_tensors(device): - """Test appending dict data to the episode.""" - dummy_dict_data_0 = { - "key_0": torch.tensor([0], device=device), - "key_1": {"key_1_0": torch.tensor([1], device=device), "key_1_1": torch.tensor([2], device=device)}, - } - dummy_dict_data_1 = { - "key_0": torch.tensor([3], device=device), - "key_1": {"key_1_0": torch.tensor([4], device=device), "key_1_1": torch.tensor([5], device=device)}, - } - - episode = EpisodeData() - - # test adding dict data to a key that does not exist - episode.add("key", dummy_dict_data_0) - key_data = episode.data.get("key") - assert key_data is not None - key_0_data = torch.stack(key_data.get("key_0")) - assert key_0_data is not None - assert torch.equal(key_0_data, torch.tensor([[0]], device=device)) - key_1_data = key_data.get("key_1") - assert key_1_data is not None - key_1_0_data = torch.stack(key_1_data.get("key_1_0")) - assert key_1_0_data is not None - assert torch.equal(key_1_0_data, torch.tensor([[1]], device=device)) - key_1_1_data = torch.stack(key_1_data.get("key_1_1")) - assert key_1_1_data is not None - assert torch.equal(key_1_1_data, torch.tensor([[2]], device=device)) - - # test adding dict data to a key that exists - episode.add("key", dummy_dict_data_1) - key_data = episode.data.get("key") - assert key_data is not None - key_0_data = torch.stack(key_data.get("key_0")) - assert key_0_data is not None - assert torch.equal(key_0_data, torch.tensor([[0], [3]], device=device)) - key_1_data = key_data.get("key_1") - assert key_1_data is not None - key_1_0_data = torch.stack(key_1_data.get("key_1_0")) - assert key_1_0_data is not None - assert torch.equal(key_1_0_data, torch.tensor([[1], [4]], device=device)) - key_1_1_data = torch.stack(key_1_data.get("key_1_1")) - assert key_1_1_data is not None - assert torch.equal(key_1_1_data, torch.tensor([[2], [5]], device=device)) - - -@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) -def test_get_initial_state(device): - """Test getting the initial state of the episode.""" - dummy_initial_state = torch.tensor([1, 2, 3], device=device) - episode = EpisodeData() - - episode.add("initial_state", dummy_initial_state) - initial_state = torch.stack(episode.get_initial_state()) - assert initial_state is not None - assert torch.equal(initial_state, dummy_initial_state.unsqueeze(0)) - - -@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) -def test_get_next_action(device): - """Test getting next actions.""" - # dummy actions - action1 = torch.tensor([1, 2, 3], device=device) - action2 = torch.tensor([4, 5, 6], device=device) - action3 = torch.tensor([7, 8, 9], device=device) - - episode = EpisodeData() - assert episode.get_next_action() is None - - episode.add("actions", action1) - episode.add("actions", action2) - episode.add("actions", action3) - - # check if actions are returned in the correct order - next_action = episode.get_next_action() - assert next_action is not None - assert torch.equal(next_action, action1) - next_action = episode.get_next_action() - assert next_action is not None - assert torch.equal(next_action, action2) - next_action = episode.get_next_action() - assert next_action is not None - assert torch.equal(next_action, action3) - - # check if None is returned when all actions are exhausted - assert episode.get_next_action() is None diff --git a/source/isaaclab/test/utils/test_hdf5_dataset_file_handler.py b/source/isaaclab/test/utils/test_hdf5_dataset_file_handler.py index aed285162e30..feb38c3ae6c4 100644 --- a/source/isaaclab/test/utils/test_hdf5_dataset_file_handler.py +++ b/source/isaaclab/test/utils/test_hdf5_dataset_file_handler.py @@ -3,10 +3,6 @@ # # SPDX-License-Identifier: BSD-3-Clause import json -import os -import shutil -import tempfile -import uuid import h5py import pytest @@ -18,69 +14,19 @@ pytestmark = pytest.mark.unit -def create_test_episode(device): - """create a test episode with dummy data.""" - test_episode = EpisodeData() - - test_episode.seed = 0 - test_episode.success = True - - test_episode.add("initial_state", torch.tensor([1, 2, 3], device=device)) - - test_episode.add("actions", torch.tensor([1, 2, 3], device=device)) - test_episode.add("actions", torch.tensor([4, 5, 6], device=device)) - test_episode.add("actions", torch.tensor([7, 8, 9], device=device)) - - test_episode.add("obs/policy/term1", torch.tensor([1, 2, 3, 4, 5], device=device)) - test_episode.add("obs/policy/term1", torch.tensor([6, 7, 8, 9, 10], device=device)) - test_episode.add("obs/policy/term1", torch.tensor([11, 12, 13, 14, 15], device=device)) - - return test_episode - - -@pytest.fixture -def temp_dir(): - """Create a temporary directory for test datasets.""" - temp_dir = tempfile.mkdtemp() - yield temp_dir - # cleanup after tests - shutil.rmtree(temp_dir) - - -def test_create_dataset_file(temp_dir, monkeypatch): - """Test creating a new dataset file.""" - # create a dataset file given a file name with extension - dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.hdf5") - dataset_file_handler = HDF5DatasetFileHandler() - dataset_file_handler.create(dataset_file_path, "test_env_name") - dataset_file_handler.close() - - # check if the dataset is created - assert os.path.exists(dataset_file_path) - - # create a dataset file given a file name without extension - dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}") - dataset_file_handler = HDF5DatasetFileHandler() - dataset_file_handler.create(dataset_file_path, "test_env_name") - dataset_file_handler.close() - - # check if the dataset is created - assert os.path.exists(dataset_file_path + ".hdf5") - - # create a dataset file given a bare file name in the current working directory - monkeypatch.chdir(temp_dir) - dataset_file_name = f"{uuid.uuid4()}.hdf5" - dataset_file_handler = HDF5DatasetFileHandler() - dataset_file_handler.create(dataset_file_name, "test_env_name") - dataset_file_handler.close() - - # check if the dataset is created - assert os.path.exists(os.path.join(temp_dir, dataset_file_name)) +def test_create_dataset_file(tmp_path, monkeypatch): + """Accept absolute paths, omitted extensions and bare filenames.""" + monkeypatch.chdir(tmp_path) + for path in (str(tmp_path / "absolute.hdf5"), str(tmp_path / "no_extension"), "relative.hdf5"): + dataset_file_handler = HDF5DatasetFileHandler() + dataset_file_handler.create(path, "test_env_name") + dataset_file_handler.close() + assert {path.name for path in tmp_path.iterdir()} == {"absolute.hdf5", "no_extension.hdf5", "relative.hdf5"} -def test_add_env_args_preserves_existing_args_after_reopen(temp_dir): - """Test extending environment arguments after reopening a dataset.""" - dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.hdf5") +def test_env_args_follow_active_file(tmp_path): + """Preserve metadata on reopen, then discard it when reusing the handler for a new file.""" + dataset_file_path = str(tmp_path / "first.hdf5") dataset_file_handler = HDF5DatasetFileHandler() dataset_file_handler.create(dataset_file_path, "test_env_name") dataset_file_handler.close() @@ -95,16 +41,7 @@ def test_add_env_args_preserves_existing_args_after_reopen(temp_dir): assert env_args == {"env_name": "test_env_name", "type": 2, "custom_arg": "custom_value"} - -def test_create_resets_env_args_when_reusing_handler(temp_dir): - """Test environment arguments do not leak between datasets.""" - first_dataset_path = os.path.join(temp_dir, f"{uuid.uuid4()}.hdf5") - second_dataset_path = os.path.join(temp_dir, f"{uuid.uuid4()}.hdf5") - dataset_file_handler = HDF5DatasetFileHandler() - dataset_file_handler.create(first_dataset_path, "first_env") - dataset_file_handler.add_env_args({"custom_arg": "custom_value"}) - dataset_file_handler.close() - + second_dataset_path = str(tmp_path / "second.hdf5") dataset_file_handler.create(second_dataset_path, "second_env") dataset_file_handler.close() @@ -114,10 +51,10 @@ def test_create_resets_env_args_when_reusing_handler(temp_dir): assert env_args == {"env_name": "second_env", "type": 2} -def test_convert_dataset_to_xyzw_updates_format_version(temp_dir): +def test_convert_dataset_to_xyzw_updates_format_version(tmp_path): """Converted datasets should load their XYZW quaternions without converting them again.""" - input_path = os.path.join(temp_dir, "legacy.hdf5") - output_path = os.path.join(temp_dir, "converted.hdf5") + input_path = str(tmp_path / "legacy.hdf5") + output_path = str(tmp_path / "converted.hdf5") with h5py.File(input_path, "w") as dataset: dataset.attrs["format_version"] = 0 data = dataset.create_group("data") @@ -140,16 +77,27 @@ def test_convert_dataset_to_xyzw_updates_format_version(temp_dir): @pytest.mark.parametrize("device", test_devices()) -def test_write_and_load_episode(temp_dir, device): - """Test writing and loading an episode to and from the dataset file.""" - dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.hdf5") +def test_write_and_load_episode(tmp_path, device): + """Round-trip episode data and metadata, then play every action through exhaustion.""" + dataset_file_path = str(tmp_path / "episodes.hdf5") dataset_file_handler = HDF5DatasetFileHandler() dataset_file_handler.create(dataset_file_path, "test_env_name") - test_episode = create_test_episode(device) + test_episode = EpisodeData() + assert test_episode.is_empty() + assert test_episode.get_next_action() is None + test_episode.seed = 0 + test_episode.success = True + test_episode.data = { + "initial_state": torch.tensor([[1, 2, 3]], device=device), + "actions": torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], device=device), + "obs": {"policy": {"term1": torch.arange(1, 16, device=device).reshape(3, 5)}}, + } + assert not test_episode.is_empty() + # Reading an absent action must not advance the cursor before data becomes available. + assert torch.equal(test_episode.get_next_action(), test_episode.data["actions"][0]) # write the episode to the dataset - test_episode.pre_export() dataset_file_handler.write_episode(test_episode) dataset_file_handler.flush() @@ -179,9 +127,12 @@ def test_write_and_load_episode(temp_dir, device): assert loaded_episode.seed == test_episode.seed assert loaded_episode.success == test_episode.success - assert torch.equal(loaded_episode.get_initial_state(), test_episode.get_initial_state()) + assert torch.equal(loaded_episode.get_initial_state(), test_episode.data["initial_state"]) + + assert torch.equal(loaded_episode.data["obs"]["policy"]["term1"], test_episode.data["obs"]["policy"]["term1"]) for action in test_episode.data["actions"]: assert torch.equal(loaded_episode.get_next_action(), action) + assert loaded_episode.get_next_action() is None dataset_file_handler.close() diff --git a/source/isaaclab/test/utils/test_math.py b/source/isaaclab/test/utils/test_math.py index d9d217931b56..1085b8f6b4b0 100644 --- a/source/isaaclab/test/utils/test_math.py +++ b/source/isaaclab/test/utils/test_math.py @@ -174,6 +174,12 @@ def test_axis_angle_from_quat_approximation(device): # Assert that the computed values are close to the expected values torch.testing.assert_close(axis_angle_computed, axis_angle_expected) + # The unselected exact formula must not introduce NaN gradients through 0/0 at identity. + identities = torch.tensor([[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.0, -1.0]], device=device, requires_grad=True) + math_utils.axis_angle_from_quat(identities).sum().backward() + expected_gradient = torch.tensor([[2.0, 2.0, 2.0, 0.0], [-2.0, -2.0, -2.0, 0.0]], device=device) + torch.testing.assert_close(identities.grad, expected_gradient) + @pytest.mark.parametrize("device", test_devices()) def test_quat_error_magnitude(device): diff --git a/source/isaaclab_contrib/changelog.d/solver-config-replacement.fixed.rst b/source/isaaclab_contrib/changelog.d/solver-config-replacement.fixed.rst new file mode 100644 index 000000000000..e2561e4261e4 --- /dev/null +++ b/source/isaaclab_contrib/changelog.d/solver-config-replacement.fixed.rst @@ -0,0 +1 @@ +* Removed the custom coupling preset workaround for copying the derived Newton manager class. diff --git a/source/isaaclab_contrib/isaaclab_contrib/custom_coupling/franka_soft_env_cfg.py b/source/isaaclab_contrib/isaaclab_contrib/custom_coupling/franka_soft_env_cfg.py index 0a112f2e733f..4059c560837f 100644 --- a/source/isaaclab_contrib/isaaclab_contrib/custom_coupling/franka_soft_env_cfg.py +++ b/source/isaaclab_contrib/isaaclab_contrib/custom_coupling/franka_soft_env_cfg.py @@ -21,8 +21,6 @@ class PhysicsCfg(CorePhysicsCfg): newton_mjwarp_vbd = replace( CorePhysicsCfg().newton_mjwarp_vbd_proxy, - # Required: ``NewtonCfg.__post_init__`` rejects a preset class_type and re-derives it. - class_type=None, solver_cfg=CoupledMJWarpVBDSolverCfg( rigid_solver_cfg=MJWarpSolverCfg( njmax=40, nconmax=20, ls_iterations=20, integrator="implicitfast", ccd_iterations=100 diff --git a/source/isaaclab_experimental/changelog.d/recorder-capture.changed.rst b/source/isaaclab_experimental/changelog.d/recorder-capture.changed.rst new file mode 100644 index 000000000000..29d7578a43f8 --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/recorder-capture.changed.rst @@ -0,0 +1,3 @@ +* **Breaking:** Removed the implicit observation computation before Warp RL post-step recorder callbacks. + Capture runs before automatic reset. Custom terms needing observations must request them explicitly + from the observation manager; ``env.obs_buf`` is no longer refreshed for recorder dispatch. diff --git a/source/isaaclab_experimental/changelog.d/recorder-capture.major b/source/isaaclab_experimental/changelog.d/recorder-capture.major new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py index be522e073a2b..5ac88958d450 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py @@ -320,14 +320,7 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: timer=DEBUG_TIMER_STEP, ) - if len(self.recorder_manager.active_terms) > 0: - # update observations for recording if needed - self._manager_call_switch.call_stage( - stage="ObservationManager_compute_no_history", - warp_call={"fn": self.observation_manager.compute, "kwargs": {"return_cloned_output": False}}, - timer=DEBUG_TIMER_STEP, - ) - self.recorder_manager.record_post_step() + self.recorder_manager.record_post_step() # -- reset envs that terminated/timed-out and log the episode information # NOTE: Interim path (intentional). diff --git a/source/isaaclab_newton/changelog.d/native-nerd.added.rst b/source/isaaclab_newton/changelog.d/native-nerd.added.rst new file mode 100644 index 000000000000..8a09b0edca71 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/native-nerd.added.rst @@ -0,0 +1,12 @@ +* Added native NeRD transition recording and a graph-enabled learned Newton solver, shared by + Cartpole and ANYmal C. Physical-model, physics-step and reset-history contracts are checked + against native training checkpoints. PhysicsNeMo Robotics is an optional experimental dependency. +* NeRD consumes only applied joint efforts, captured after actuator evaluation and clipping and before + the solver runs. Capture/deployment guards reject solver-side motor drives and external body wrenches. + Physical identity retains fixed-root geometry, body flags and collision filtering; policy decimation + repeats the actuator and learned physics step. Target-conditioned models are rejected. +* Added shared live Newton contact encoding for rough-terrain recording and learned deployment, + with synchronized reset histories, explicit capacity checks and fixed contact-material identity. + +* Delegated heading conversion, temporal cache/reset handling and learned Newton state writeback to + PhysicsNeMo Robotics. ParallelNeRD deployment uses resident streaming inside the host CUDA graph. diff --git a/source/isaaclab_newton/changelog.d/native-nerd.minor b/source/isaaclab_newton/changelog.d/native-nerd.minor new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_newton/changelog.d/solver-config-replacement.fixed.rst b/source/isaaclab_newton/changelog.d/solver-config-replacement.fixed.rst new file mode 100644 index 000000000000..3657beafee84 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/solver-config-replacement.fixed.rst @@ -0,0 +1,2 @@ +* Made ``NewtonCfg.class_type`` derived constructor state so replacing ``solver_cfg`` selects + the new manager without copying a stale manager class into the configuration constructor. diff --git a/source/isaaclab_newton/isaaclab_newton/envs/mdp/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/envs/mdp/__init__.pyi index a0b4c9f49d3d..4ee2bb50b321 100644 --- a/source/isaaclab_newton/isaaclab_newton/envs/mdp/__init__.pyi +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/__init__.pyi @@ -14,6 +14,9 @@ __all__ = [ "NewtonOperationalSpaceControllerAction", "NewtonOperationalSpaceControllerActionCfg", "randomize_visual_shape", + "NeRDRecorder", + "NeRDRecorderCfg", + "NeRDRecorderManagerCfg", ] from .actions import ( @@ -30,3 +33,4 @@ from .events import ( randomize_rigid_body_material, randomize_visual_shape, ) +from .recorders import NeRDRecorder, NeRDRecorderCfg, NeRDRecorderManagerCfg diff --git a/source/isaaclab_newton/isaaclab_newton/envs/mdp/recorders.py b/source/isaaclab_newton/isaaclab_newton/envs/mdp/recorders.py new file mode 100644 index 000000000000..99ccb8c5fdc7 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/recorders.py @@ -0,0 +1,112 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Native recorder terms for generalized-coordinate learned dynamics.""" + +from contextlib import nullcontext + +import torch +import warp as wp + +from isaaclab.managers import DatasetExportMode, RecorderManagerBaseCfg, RecorderTerm, RecorderTermCfg +from isaaclab.utils import configclass + + +@configclass +class NeRDRecorderCfg(RecorderTermCfg): + """Select the transition representation independently of its consumer or storage.""" + + class_type: str = "{DIR}.recorders:NeRDRecorder" + heading: bool = False + contact_bodies: tuple[str, ...] = () + root_body: str = "" + max_contacts: int = 0 + """Positive capacity selects live Newton contact sets instead of fixed plane probes.""" + record_to_episodes: bool = False + """Retain transitions through RecorderManager; streaming consumers read ``latest`` either way.""" + + +@configclass +class NeRDRecorderManagerCfg(RecorderManagerBaseCfg): + """Optional native dynamics capture; task playback has no capture cost by default.""" + + dataset_export_mode: DatasetExportMode = DatasetExportMode.EXPORT_NONE + transition: NeRDRecorderCfg | None = None + + +class NeRDRecorder(RecorderTerm): + """Capture the terminal successor before autoreset into persistent device tensors. + + ``latest`` is valid until the next step. Replay and curiosity consumers copy the fields they + retain; this term owns neither a dataset writer nor an episode/replay buffer. + """ + + def __init__(self, cfg, env): + super().__init__(cfg, env) + from isaaclab_newton.physics.nerd import JointDynamics + + self.manager = env.sim.backend_manager + if env.cfg.decimation != 1 or env.cfg.sim.physics.num_substeps != 1: + raise ValueError("Effort recording requires one physics step per environment step.") + model = self.manager.get_model() + self.dynamics = JointDynamics( + model=model, + frame_dt=env.step_dt, + heading=cfg.heading, + contact_bodies=cfg.contact_bodies, + root_body=cfg.root_body, + max_contacts=cfg.max_contacts, + ) + self.control = self.manager.get_control() + self.contacts = self.manager.get_contacts() + self._torch_pool = torch.cuda.MemPool() if cfg.max_contacts and model.device.is_cuda else None + if cfg.max_contacts and not self.manager._needs_collision_pipeline: + raise ValueError("Newton contact capture requires the Newton collision pipeline.") + self.manager.register_state_force_callback(self._capture_solver_inputs) + width = self.dynamics.codec.state_shape[0] + self.latest = { + "state": torch.empty(self.num_envs, width, device=self.device), + "next_state": torch.empty(self.num_envs, width, device=self.device), + "control": torch.empty(self.num_envs, self.dynamics.control_dim, device=self.device), + "contacts": torch.empty(self.num_envs, *self.dynamics.contact_shape, device=self.device), + "terminated": torch.empty(self.num_envs, dtype=torch.bool, device=self.device), + "truncated": torch.empty(self.num_envs, dtype=torch.bool, device=self.device), + } + + self._has_transition = False + + def record_pre_step(self): + state = self.manager.get_scene_data_backend().state + self.dynamics.codec.read_into(state, self.latest["state"]) + if not self.cfg.max_contacts: + self.latest["contacts"].copy_(self.dynamics.contacts(state)) + return None, None + + def _capture_solver_inputs(self, state): + device = state.joint_q.device + stream = torch.cuda.stream(wp.stream_to_torch(wp.get_stream(device))) if device.is_cuda else nullcontext() + pool = ( + torch.cuda.use_mem_pool(self._torch_pool, device=device.alias) + if self._torch_pool is not None + else nullcontext() + ) + with stream, pool: + self.dynamics.check_inputs(state) + self.latest["control"].copy_(wp.to_torch(self.control.joint_f).reshape(self.num_envs, -1)) + if self.cfg.max_contacts: + self.latest["contacts"].copy_(self.dynamics.contacts(state, self.contacts)) + + def record_post_step(self): + self.dynamics.codec.read_into(self.manager.get_state_0(), self.latest["next_state"]) + self.latest["terminated"].copy_(self._env.reset_terminated) + self.latest["truncated"].copy_(self._env.reset_time_outs) + self._has_transition = True + return ("transition", self.latest) if self.cfg.record_to_episodes else (None, None) + + def record_pre_reset(self, env_ids): + if self._has_transition: + ids = slice(None) if env_ids is None else env_ids + self.latest["truncated"][ids] |= ~self.latest["terminated"][ids] + return None, None diff --git a/source/isaaclab_newton/isaaclab_newton/physics/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/physics/__init__.pyi index 4a18a1fcc015..ad6f2c43fbaf 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/__init__.pyi +++ b/source/isaaclab_newton/isaaclab_newton/physics/__init__.pyi @@ -17,6 +17,7 @@ __all__ = [ "KaminoPADMMSolverCfg", "MPMSolverCfg", "MJWarpSolverCfg", + "NeRDSolverCfg", "MjWarpActuatorBridge", "NewtonCfg", "NewtonBackendCfg", @@ -28,6 +29,7 @@ __all__ = [ "NewtonManager", "NewtonQueries", "NewtonMJWarpManager", + "NewtonNeRDManager", "NewtonShapeCfg", "NewtonSoftContactCfg", "NewtonSolverCfg", @@ -57,6 +59,8 @@ from .mjwarp_manager import NewtonMJWarpManager from .mjwarp_manager_cfg import MJWarpSolverCfg from .mpm_manager import NewtonMPMManager from .mpm_manager_cfg import MPMSolverCfg +from .nerd_manager import NewtonNeRDManager +from .nerd_manager_cfg import NeRDSolverCfg from .newton_collision_cfg import HydroelasticSDFCfg, NewtonCollisionPipelineCfg from .newton_manager import NewtonManager, NewtonQueries, create_newton_builder from .newton_manager_cfg import ( diff --git a/source/isaaclab_newton/isaaclab_newton/physics/nerd.py b/source/isaaclab_newton/isaaclab_newton/physics/nerd.py new file mode 100644 index 000000000000..cdff59be54e4 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/physics/nerd.py @@ -0,0 +1,412 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Shared NeRD representation for native recording, offline fitting and Newton inference. + +PhysicsNeMo Robotics owns state conversion, contact encoders and temporal inference. Isaac Lab +owns the recorded physics-step contract, physical identity policy and task feature selection. +Import this module only when learned dynamics is selected; task discovery needs no PhysicsNeMo. +""" + +from __future__ import annotations + +import copy +import hashlib +import json +from contextlib import nullcontext + +import newton +import numpy as np +import torch +import warp as wp +from physicsnemo_robotics.dynamics import ( + BodyHeadingFrame, + BodyStateCodec, + CodecSignature, + FixedContactFeatureAdapter, + FixedRigidContactEncoder, + GravityEncoder, + JointStateCodec, + RigidContactSetEncoder, + expand_rigid_contact_shapes, +) +from physicsnemo_robotics.dynamics.codecs.joint import JointLayout +from physicsnemo_robotics.newton import NewtonStepContext +from tensordict import TensorDict + + +class JointDynamics: + """One representation used on stored tensors and, optionally, a live Newton model. + + Applied joint efforts are the sole control input; actuators run outside learned physics. + Only uniform, joint-backed rigid worlds are supported. Contact inputs are fixed plane probes + or live Newton rigid contact sets, recomputed from each predicted state during deployment. + The descriptor is JSON data, so offline training does not initialize a simulation. + """ + + def __init__( + self, + descriptor=None, + model=None, + *, + frame_dt=None, + heading=False, + contact_bodies=(), + root_body="", + max_contacts=0, + ): + if descriptor is None: + if max_contacts < 0 or (max_contacts and contact_bodies): + raise ValueError("Select fixed contact_bodies or a positive Newton max_contacts capacity.") + if ( + model.particle_count + or model.body_count != model.joint_count + or model.articulation_count != model.world_count + ): + raise ValueError("NeRD requires one complete joint-backed rigid articulation per world.") + self.codec = JointStateCodec(model, heading=heading) + descriptor = { + "version": 2, + "frame_dt": frame_dt, + "control": "joint_f", + "heading": heading, + "contact_bodies": list(contact_bodies), + "root_body": root_body, + "max_contacts": max_contacts, + "contact_entities": [ + model.body_label[i].rsplit("/", 1)[-1] for i in np.flatnonzero(model.body_world.numpy() == 0) + ] + if max_contacts + else [], + "codec": self.codec.signature.to_dict(), + "model": model_signature(model, contact_conditioned=bool(max_contacts)), + } + else: + if descriptor["version"] != 2: + raise ValueError("NeRD requires descriptor version 2; recollect data with complete physical identity.") + if descriptor["control"] != "joint_f": + raise ValueError("NeRD requires applied joint efforts; recollect and refit target-conditioned data.") + if model is not None: + self.codec = JointStateCodec(model, heading=descriptor["heading"]) + elif "codec" in descriptor: + self.codec = JointStateCodec.from_signature(CodecSignature.from_dict(descriptor["codec"])) + else: + # Existing raw v2 recordings precede portable codec signatures. Their + # joint layout still describes the stored world-space state exactly. + layout = dict(descriptor["layout"]) + for name in ("continuous_q_mask", "root_translation_mask"): + layout[name] = torch.tensor(layout[name], dtype=torch.bool) + for name in ("quaternion_q_starts", "free_root_translation_q_starts"): + layout[name] = tuple(layout[name]) + layout["world_count"] = 1 + self.codec = JointStateCodec(layout=JointLayout(**layout), heading=descriptor["heading"]) + if self.codec.heading != descriptor["heading"] or ( + "codec" in descriptor + and self.codec.signature.fingerprint != CodecSignature.from_dict(descriptor["codec"]).fingerprint + ): + raise ValueError("NeRD codec signature differs from the recorded state/heading convention.") + self.descriptor = descriptor + self.gravity = GravityEncoder(self.codec, (0.0, 0.0, -1.0)) if descriptor["heading"] else None + self.max_contacts = descriptor.get("max_contacts", 0) + if model is not None: + if ( + np.any(model.joint_target_ke.numpy()) + or np.any(model.joint_target_kd.numpy()) + or model.custom_frequency_counts.get("mujoco:actuator", 0) + ): + raise ValueError( + "NeRD requires all motor effort in joint_f; solver-side drives are unsupported. " + "Compute and clip efforts in the actuator before each physics step." + ) + self._input_valid = torch.empty((), dtype=torch.bool, device=model.device.alias) + self.contact_encoder = self.contact_adapter = None + if model is not None and self.max_contacts: + body_worlds = model.body_world.numpy() + indices = np.stack([np.flatnonzero(body_worlds == i) for i in range(model.world_count)]) + names = [model.body_label[i].rsplit("/", 1)[-1] for i in indices[0]] + if names != descriptor["contact_entities"]: + raise ValueError("NeRD contact body order differs from the recorded representation.") + body_codec = BodyStateCodec( + model, + indices=indices, + semantic_ids=tuple(names), + model_frame=BodyHeadingFrame(body=names.index(descriptor["root_body"])) + if descriptor["heading"] + else None, + ) + self.contact_encoder = RigidContactSetEncoder(model, body_codec, max_contacts=self.max_contacts) + elif model is not None and descriptor["contact_bodies"]: + # The probe library compares local geometry across replicas. USD world-transform + # roundoff is removed from a private view; live model arrays are never replaced. + geometry = copy.copy(model) + transforms = model.shape_transform.numpy().astype(np.float64) + transforms[np.abs(transforms) < 1.0e-8] = 0.0 + geometry.shape_transform = wp.array(np.round(transforms, 6), dtype=wp.transform, device=model.device) + body_ids = np.flatnonzero(model.body_world.numpy() == 0) + body_names = [model.body_label[i].rsplit("/", 1)[-1] for i in body_ids] + shape_ids = np.flatnonzero(model.shape_world.numpy() == 0) + shape_body, shape_type = model.shape_body.numpy(), model.shape_type.numpy() + feet = [] + for name in descriptor["contact_bodies"]: + body = body_ids[body_names.index(name)] + slots = [ + i + for i, s in enumerate(shape_ids) + if shape_body[s] == body and shape_type[s] == newton.GeoType.SPHERE + ] + if len(slots) != 1: + raise ValueError(f"NeRD contact body {name!r} needs exactly one sphere.") + feet.append(slots[0]) + planes = [ + i for i, s in enumerate(shape_ids) if shape_body[s] == -1 and shape_type[s] == newton.GeoType.PLANE + ] + if len(planes) != 1: + raise ValueError("NeRD fixed contacts need exactly one plane in each world.") + body_codec = BodyStateCodec( + geometry, model_frame=BodyHeadingFrame(body=body_names.index(descriptor["root_body"])) + ) + self.contact_encoder = FixedRigidContactEncoder( + geometry, + body_codec, + points=expand_rigid_contact_shapes(geometry, body_codec, feet), + environment_shape_slots=planes, + environment_geometry="per_world", + key="contact/probes", + ) + self.contact_adapter = FixedContactFeatureAdapter( + self.contact_encoder, ("point", "normal", "lever_arm", "gap", "relative_velocity", "contact_state") + ) + + @property + def control_dim(self): + return self.codec.layout.dof_qd + + @property + def contact_dim(self): + return 14 * len(self.descriptor["contact_bodies"]) + + @property + def contact_shape(self): + return (self.max_contacts, 17) if self.max_contacts else (self.contact_dim,) + + @property + def contact_model_cfg(self): + if not self.max_contacts: + return {} + return dict( + contact_dim=17, + max_contacts=self.max_contacts, + num_contact_entities=len(self.descriptor["contact_entities"]), + ) + + @property + def input_dim(self): + return ( + self.codec.state_shape[0] + self.control_dim + (3 if self.descriptor["heading"] else 0) + self.contact_dim + ) + + def inputs(self, state, control, contacts): + """Encode one frame or a batch of sequences without discarding autograd.""" + values = [self.codec.encode_state(state), control] + if self.descriptor["heading"]: + values.append(self.gravity.encode_window(TensorDict({"history/state": state}, batch_size=[]))) + if not self.max_contacts: + values.append(contacts) + return torch.cat(values, dim=-1) + + def contacts(self, state, contacts=None): + """Encode starting-state geometry; variable contacts must come from the live collision pipeline.""" + if self.max_contacts: + dt = self.descriptor["frame_dt"] + context = NewtonStepContext(state, state, None, contacts, 0, 0, dt, dt) + values = self.contact_encoder.encode_step(context) + torch._assert_async((self.contact_encoder.last_overflow == 0).all(), "NeRD contact capacity overflowed.") + return values + if self.contact_encoder is None: + return torch.empty(self.codec.batch_size, 0, device=state.joint_q.device.alias) + return self.contact_adapter.encode_tokens(self.contact_encoder.encode_state(state)) + + def check_inputs(self, state): + """Reject unrepresented inputs at the same solver boundary during capture and deployment.""" + # Persistent reduction outputs also work inside the numerical teacher's Warp CUDA graph. + device = state.body_f.device + stream = torch.cuda.stream(wp.stream_to_torch(wp.get_stream(device))) if device.is_cuda else nullcontext() + with stream: + torch.any(wp.to_torch(state.body_f), out=self._input_valid) + torch._assert_async(self._input_valid.logical_not_(), "NeRD has no external body-wrench input.") + + +def model_signature(model: newton.Model, *, contact_conditioned: bool = False) -> str: + """Fingerprint the supported rigid model without mutating it or including solver/viewer attributes. + + Replicated worlds must have the same physical parameters. Local transform roundoff from USD + placement is tolerated at 1e-6; masses, gains, limits and friction remain exact. Horizontal + translation is ignored only in scenes with no static geometry except horizontal planes. + Fixed-root orientation and height remain physical parameters. Mesh geometry is included by content. + Live-contact models instead fix robot geometry and uniform external material, allowing static + terrain geometry and floating-root placement to vary through the encoded state and contacts. + """ + digest = hashlib.sha256() + collision_shapes = (model.shape_flags.numpy() & int(newton.ShapeFlags.COLLIDE_SHAPES)) != 0 + shape_types, shape_bodies = model.shape_type.numpy(), model.shape_body.numpy() + transforms = model.shape_transform.numpy() + static = collision_shapes & (shape_bodies == -1) + horizontal = (shape_types == newton.GeoType.PLANE) & np.all(np.abs(transforms[:, 3:5]) < 1.0e-8, axis=1) + translation_invariant = np.all(horizontal[static]) + if contact_conditioned: + # Geometry is a live input. Its contact law is not: every external surface must share it. + for name in ( + "margin", + "gap", + "material_ke", + "material_kd", + "material_mu", + "material_kf", + "material_ka", + "material_restitution", + "material_mu_torsional", + "material_mu_rolling", + ): + values = getattr(model, "shape_" + name).numpy()[static] + if not len(values) or not np.all(values == values[0]): + raise ValueError("Contact-conditioned NeRD requires one uniform external contact material.") + digest.update(name.encode()) + digest.update(np.asarray(values[0]).tobytes()) + collision_shapes &= ~static + groups = { + "body": ("mass", "inertia", "com", "flags"), + "joint": ("type", "dof_dim", "enabled", "X_p", "X_c"), + "shape": ( + "type", + "scale", + "transform", + "margin", + "gap", + "material_ke", + "material_kd", + "material_mu", + "material_kf", + "material_ka", + "material_restitution", + "material_mu_torsional", + "material_mu_rolling", + ), + } + for entity, names in groups.items(): + worlds = getattr(model, entity + "_world").numpy().copy() + if entity == "shape": + # Rendering can add world-positioned visual meshes. They do not change rigid dynamics. + worlds[~collision_shapes] = -2 + counts = np.bincount(worlds[worlds >= 0], minlength=model.world_count) + if np.any(counts != counts[0]): + raise ValueError(f"NeRD needs identical {entity} counts in each world.") + local = np.flatnonzero(worlds >= 0) + ids = local[np.argsort(worlds[local], kind="stable")].reshape(model.world_count, -1) + labels = getattr(model, entity + "_label") + digest.update(json.dumps([labels[i].split("/env_0/")[-1] for i in ids[0]]).encode()) + for suffix in names: + name = entity + "_" + suffix + array = getattr(model, name).numpy().copy() + if name == "joint_X_p" and translation_invariant: + array[model.joint_parent.numpy() == -1, :2] = 0.0 + if name == "joint_X_p" and contact_conditioned: + roots = (model.joint_parent.numpy() == -1) & (model.joint_type.numpy() == newton.JointType.FREE) + array[roots, :3] = 0.0 + if name == "shape_transform": + # A horizontal plane is unchanged by translation in its own plane. + array[horizontal, :2] = 0.0 + if suffix in ("X_p", "X_c", "transform"): + array = np.round(array.astype(np.float64), 6) + array[np.abs(array) < 1.0e-8] = 0.0 + first = array[ids[0]] + if not np.all(array[ids] == first): + raise ValueError(f"NeRD cannot represent per-world variation in {name}.") + digest.update(name.encode()) + digest.update(first.tobytes()) + digest.update(array[worlds == -1].tobytes()) + for name in ( + "joint_axis", + "joint_limit_lower", + "joint_limit_upper", + "joint_armature", + "joint_friction", + "joint_target_ke", + "joint_target_kd", + "joint_damping", + "joint_effort_limit", + "joint_velocity_limit", + "joint_limit_ke", + "joint_limit_kd", + ): + values = getattr(model, name).numpy().reshape(model.world_count, -1) + if not np.all(values == values[0]): + raise ValueError(f"NeRD cannot represent per-world variation in {name}.") + digest.update(name.encode()) + digest.update(values[0].tobytes()) + body_worlds = model.body_world.numpy() + body_ids = np.argsort(body_worlds, kind="stable").reshape(model.world_count, -1) + body_slots = np.empty(model.body_count, dtype=np.int64) + body_slots[body_ids] = np.arange(body_ids.shape[1]) + for name, values in ( + ("joint_parent", model.joint_parent.numpy()), + ("joint_child", model.joint_child.numpy()), + ("shape_body", model.shape_body.numpy()), + ): + entity = "shape" if name == "shape_body" else "joint" + worlds = getattr(model, entity + "_world").numpy() + selected = worlds >= 0 + if entity == "shape": + selected &= collision_shapes + indices = np.flatnonzero(selected) + indices = indices[np.argsort(worlds[indices], kind="stable")].reshape(model.world_count, -1) + values = np.where(values >= 0, body_slots[values.clip(min=0)], -1) + local = values[indices] + if not np.all(local == local[0]): + raise ValueError(f"NeRD cannot represent per-world variation in {name}.") + digest.update(name.encode()) + digest.update(local[0].tobytes()) + for shape in np.flatnonzero((model.shape_world.numpy() == 0) & collision_shapes): + source = model.shape_source[shape] + if source is not None: + if not isinstance(source, newton.Mesh): + raise ValueError("NeRD supports primitive and mesh collision geometry only.") + digest.update(np.asarray(source.vertices).tobytes()) + digest.update(np.asarray(source.indices).tobytes()) + # Pair indices include visual shapes and all replicas. Normalize them to collision-only + # slots so headless collection and playback with a different world count have the same identity. + shape_world = model.shape_world.numpy() + local_shapes = np.flatnonzero(collision_shapes & (shape_world >= 0)) + local_shapes = local_shapes[np.argsort(shape_world[local_shapes], kind="stable")].reshape(model.world_count, -1) + slots = np.full(model.shape_count, -1, dtype=np.int64) + slots[local_shapes] = np.arange(local_shapes.shape[1]) + global_shapes = np.flatnonzero(collision_shapes & (shape_world == -1)) + slots[global_shapes] = -1 - np.arange(len(global_shapes)) + for name, pairs in ( + ("excluded_pairs", model.shape_collision_filter_pairs_array()), + ("contact_pairs", model.shape_contact_pairs.numpy()), + ): + pairs = np.asarray(pairs).reshape(-1, 2) + pairs = pairs[collision_shapes[pairs].all(axis=1)] + pair_worlds = shape_world[pairs] + # Different local worlds cannot collide; a global shape can participate in each world. + pairs = pairs[(pair_worlds[:, 0] == pair_worlds[:, 1]) | (pair_worlds == -1).any(axis=1)] + worlds = shape_world[pairs].max(axis=1) + canonical = np.sort(slots[pairs], axis=1) + order = np.lexsort((canonical[:, 1], canonical[:, 0], worlds)) + canonical, worlds = canonical[order], worlds[order] + counts = np.bincount(worlds[worlds >= 0], minlength=model.world_count) + if np.any(counts != counts[0]): + raise ValueError(f"NeRD cannot represent per-world variation in {name}.") + local = canonical[worlds >= 0].reshape(model.world_count, -1, 2) + if not np.all(local == local[0]): + raise ValueError(f"NeRD cannot represent per-world variation in {name}.") + digest.update(name.encode()) + digest.update(local[0].tobytes()) + digest.update(canonical[worlds == -1].tobytes()) + gravity = model.gravity.numpy() + if not np.all(gravity == gravity[0]): + raise ValueError("NeRD cannot represent per-world gravity variation.") + digest.update(np.asarray(gravity[0]).tobytes()) + return digest.hexdigest() diff --git a/source/isaaclab_newton/isaaclab_newton/physics/nerd_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/nerd_manager.py new file mode 100644 index 000000000000..ef7735d45cac --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/physics/nerd_manager.py @@ -0,0 +1,198 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Native Newton solver boundary for a temporal NeRD model.""" + +import math + +import newton +import torch +import warp as wp +from physicsnemo_robotics.dynamics import ( + CompositeFeatureEncoder, + ControlFieldEncoder, + FieldNormalizer, + LearnedDynamicsSolver, + TrajectoryFieldSpec, + TrajectorySchema, + TransitionAdapter, + TransitionSpec, +) +from physicsnemo_robotics.newton import ( + NewtonExecutionSpec, + NewtonModelSignature, + NewtonSimulationSpec, + NewtonSolverSignature, + NewtonTimeSpec, +) +from physicsnemo_robotics.newton.simulation import collision_pipeline_configuration + +from isaaclab.utils import string_to_callable + +from .nerd import JointDynamics, model_signature +from .newton_manager import NewtonManager + + +class SolverNeRD(LearnedDynamicsSolver): + """Native checkpoint and effort/identity policy around PhysicsNeMo's resident solver.""" + + def __init__(self, model, checkpoint, *, compile=True, temporal="auto", collision_args=None): + if not model.device.is_cuda: + raise ValueError("The native NeRD solver currently requires CUDA.") + saved = torch.load(checkpoint, map_location="cpu", weights_only=True) + if saved["format"] != "isaaclab-nerd-v3": + raise ValueError( + "Refit the recorded transitions to a v3 checkpoint with canonical PhysicsNeMo heading inputs." + ) + self.dynamics = JointDynamics(saved["dynamics"], model) + if model_signature(model, contact_conditioned=bool(self.dynamics.max_contacts)) != saved["dynamics"]["model"]: + raise ValueError("NeRD checkpoint physical model differs from the live scene.") + network = ( + string_to_callable(saved["network_type"])( + self.dynamics.input_dim, + self.dynamics.codec.prediction_shape[0], + **saved["network_cfg"], + **self.dynamics.contact_model_cfg, + ) + .to(model.device.alias) + .eval() + ) + network.load_state_dict(saved["model"]) + network.requires_grad_(False) + history = 1 + sum(span - 1 for span in saved["network_cfg"]["temporal_attention_spans"]) + collision = newton.CollisionPipeline(model, **(collision_args or {})) if self.dynamics.max_contacts else None + # This is the live deployment contract. The native checkpoint identity was + # checked above; no teacher solver or portable artifact is inferred here. + features = [ControlFieldEncoder("joint_f", per_world_shape=(self.dynamics.control_dim,))] + fields = [ + TrajectoryFieldSpec( + "control/joint_f", "transition", dtype="float32", trailing_shape=(self.dynamics.control_dim,) + ) + ] + if self.dynamics.gravity is not None: + features.append(self.dynamics.gravity) + if self.dynamics.contact_adapter is not None: + features.append(self.dynamics.contact_adapter) + fields.append(self.dynamics.contact_adapter.field_spec) + if self.dynamics.max_contacts: + fields.append(self.dynamics.contact_encoder.field_spec) + identity = NewtonModelSignature( + topology="isaaclab-uniform-joint-worlds", + static_fingerprint=self.dynamics.descriptor["model"], + newton_version=newton.__version__, + ) + simulation = NewtonSimulationSpec( + model=identity, + solver=NewtonSolverSignature("isaaclab_newton.physics.nerd_manager.SolverNeRD"), + time=NewtonTimeSpec(solver_dt=self.dynamics.descriptor["frame_dt"], substeps=1), + ) + execution = NewtonExecutionSpec( + stepper_implementation="isaaclab_newton.physics.nerd_manager.NewtonNeRDManager", + collision_pipeline_implementation=( + f"{type(collision).__module__}.{type(collision).__qualname__}" if collision is not None else None + ), + resolved_collision_configuration=( + collision_pipeline_configuration(collision) if collision is not None else {} + ), + collide_on_reset=False, + collide_each_substep=bool(self.dynamics.max_contacts), + clear_forces=True, + requires_grad=False, + ) + schema = TrajectorySchema.from_specs(simulation, execution, self.dynamics.codec, fields=fields) + adapter = TransitionAdapter( + schema, + self.dynamics.codec, + TransitionSpec(history=history), + features=CompositeFeatureEncoder(*features), + contact_features=self.dynamics.contact_encoder if self.dynamics.max_contacts else None, + ) + + # Only models promising repeated-frame per-world resets can use cold streaming. + if temporal == "auto": + temporal = "stream" if getattr(network, "supports_repeated_frame_reset", False) else "full_window" + super().__init__( + model, + network, + adapter, + normalizer=FieldNormalizer.from_dict(saved["normalizer"]), + input_normalizer_key="input", + prediction_normalizer_key="target", + contact_normalizer_key="contact" if self.dynamics.max_contacts else None, + model_signature=adapter.schema.simulation.model, + continuity="off", + temporal=temporal, + compile=compile, + ) + # Compile once on disposable state; the manager later captures the entire physics frame. + state, control = model.state(), model.control() + contacts = None + if collision is not None: + contacts = collision.contacts() + collision.collide(state, contacts) + with torch.no_grad(): + self.step(state, state, control, contacts, self.frame_dt) + self.reset(state) + + @torch.no_grad() + def step(self, state_in, state_out, control, contacts, dt, *, observables=None): + self.dynamics.check_inputs(state_in) + super().step(state_in, state_out, control, contacts, dt, observables=observables) + if self.dynamics.max_contacts: + with torch.cuda.stream(wp.stream_to_torch(wp.get_stream(self.model.device))): + torch._assert_async( + (self.dynamics.contact_encoder.last_overflow == 0).all(), "NeRD contact capacity overflowed." + ) + + def notify_model_changed(self, flags): + # Fixed-base root-pose resets can notify unchanged joint properties. Verify the same + # physical contract instead of suppressing every notification or mutating a library method. + if ( + int(flags) + and model_signature(self.model, contact_conditioned=bool(self.dynamics.max_contacts)) + != self.dynamics.descriptor["model"] + ): + raise ValueError("Physical model changes invalidate the trained NeRD checkpoint.") + + +class NewtonNeRDManager(NewtonManager): + """Construct learned physics through the same instance lifecycle as other Newton solvers.""" + + def _build_solver(self, model, solver_cfg): + if self._num_substeps != 1 or not solver_cfg.checkpoint: + raise ValueError("NeRD requires num_substeps=1 and a checkpoint path.") + if ( + set(model.get_requested_contact_attributes()) & newton.Contacts.EXTENDED_ATTRIBUTES + or set(model.get_requested_state_attributes()) & newton.State.EXTENDED_ATTRIBUTES + ): + raise ValueError("NeRD does not produce extended sensor state or solved contact forces.") + self._solver = SolverNeRD( + model, + solver_cfg.checkpoint, + compile=solver_cfg.compile, + temporal=solver_cfg.temporal, + collision_args=self._collision_cfg.to_pipeline_args() if self._collision_cfg is not None else None, + ) + if not math.isclose(self._solver_dt, self._solver.frame_dt, rel_tol=1e-9): + raise ValueError("Set sim.dt to the NeRD checkpoint's physics-step duration.") + self._use_single_state = True + self._needs_collision_pipeline = bool(self._solver.dynamics.max_contacts) + self._supports_rigid_body_force_input = False + self._supports_contact_sensors = False + + def _capture_graph(self, capture_target): + # Warp owns graph capture/replay. Isolate its Torch temporaries from the caller's caching + # allocator and keep this pool alive as long as the captured physics graph can use them. + self._torch_pool = torch.cuda.MemPool() + with torch.cuda.use_mem_pool(self._torch_pool, device=self._device): + return super()._capture_graph(capture_target) + + def clear(self): + if self._solver is not None: + self._solver.session.close() + super().clear() + + def _solver_specific_clear(self): + self._torch_pool = None diff --git a/source/isaaclab_newton/isaaclab_newton/physics/nerd_manager_cfg.py b/source/isaaclab_newton/isaaclab_newton/physics/nerd_manager_cfg.py new file mode 100644 index 000000000000..db8e2313dc7f --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/physics/nerd_manager_cfg.py @@ -0,0 +1,38 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configuration of learned generalized-coordinate Newton dynamics.""" + +from typing import Literal + +from isaaclab.utils import configclass + +from .newton_manager_cfg import NewtonSolverCfg + + +@configclass +class NeRDSolverCfg(NewtonSolverCfg): + """Advance one physics step from applied joint efforts with a native NeRD checkpoint. + + Use one physics substep. Policy decimation repeats actuator evaluation and learned physics; + collection records each physics step separately. The checkpoint declares the frame duration, + modeled joints and contact features. Actuator commands, state and effort limits remain owned + by the actuator. Solver-side motor drives are unsupported. + """ + + class_type: str = "{DIR}.nerd_manager:NewtonNeRDManager" + solver_type: str = "nerd" + checkpoint: str = "" + """Native checkpoint written by ``scripts/learning/nerd/train.py``.""" + + compile: bool = True + """Compile model inference before the native manager captures its physics CUDA graph.""" + + temporal: Literal["auto", "full_window", "stream"] = "auto" + """PhysicsNeMo temporal execution; auto uses caching when selected-world resets are supported. + + Both paths run inside the manager's physics CUDA graph. Full-window execution is useful for + numerical comparison; requesting stream requires a compatible network reset contract. + """ diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py index 18687ef9a063..88efcae3024c 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py @@ -161,11 +161,11 @@ class NewtonCfg(PhysicsCfg): :attr:`class_type` explicitly. """ - class_type: type[NewtonManager] | str | None = None + class_type: type[NewtonManager] | str | None = field(default=None, init=False) """The class type of the :class:`NewtonManager`. Auto-set in :meth:`__post_init__` from :attr:`solver_cfg.class_type`. - Users normally do not set this directly. + It is not a constructor argument, so replacing a solver derives the matching manager again. """ num_substeps: int = 1 @@ -294,10 +294,6 @@ class NewtonCfg(PhysicsCfg): def __post_init__(self): # NewtonCfg.class_type is auto-derived from solver_cfg.class_type. - # Refuse a user-set value: setting both is ambiguous and was - # previously silently overwritten. - if self.class_type is not None: - raise TypeError("Cannot manually set NewtonCfg.class_type; it is auto-derived from solver_cfg.class_type.") if self.deterministic_mode not in ("not_guaranteed", "run_to_run", "gpu_to_gpu"): raise ValueError( "NewtonCfg.deterministic_mode must be 'not_guaranteed', 'run_to_run', or 'gpu_to_gpu', " diff --git a/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py b/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py index 51c354e87c47..de0f0a019eea 100644 --- a/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py +++ b/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py @@ -167,6 +167,8 @@ def test_newton_cfg_post_init_propagates_class_type( """``NewtonCfg.__post_init__`` lifts ``solver_cfg.class_type`` onto ``NewtonCfg.class_type``.""" cfg = NewtonCfg(solver_cfg=solver_cfg_factory()) assert cfg.class_type.__name__ == expected_manager.__name__ + replaced = NewtonCfg().replace(solver_cfg=solver_cfg_factory()) + assert replaced.class_type.__name__ == expected_manager.__name__ @pytest.mark.parametrize( diff --git a/source/isaaclab_newton/test/test_nerd.py b/source/isaaclab_newton/test/test_nerd.py new file mode 100644 index 000000000000..122b9b7d0e5e --- /dev/null +++ b/source/isaaclab_newton/test/test_nerd.py @@ -0,0 +1,476 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""NeRD capture, fitting, checkpoint and graph/reset contracts across articulated tasks and contact models.""" + +import ast +import json +import subprocess +import sys +from pathlib import Path + +import gymnasium as gym +import h5py +import newton +import pytest +import torch +import warp as wp + +from isaaclab.app import launch_simulation +from isaaclab.managers import DatasetExportMode +from isaaclab.paths import ISAACLAB_ROOT +from isaaclab.test.utils import DeviceScope, test_devices +from isaaclab.utils import to_dict + +from isaaclab_tasks.utils import resolve_task_config + + +def test_native_nerd_ownership(): + """Keep capture, solver state and learning dependencies at their owning boundaries.""" + source = Path(ISAACLAB_ROOT) / "source" + assert not (source / "isaaclab_physicsnemo").exists() + newton = source / "isaaclab_newton/isaaclab_newton" + task_core = source / "isaaclab_tasks/isaaclab_tasks/core" + paths = ( + *newton.glob("physics/nerd*.py"), + newton / "envs/mdp/recorders.py", + task_core / "cartpole/cartpole_nerd_env_cfg.py", + *task_core.glob("velocity/config/anymal_c/*.py"), + source / "isaaclab_rl/isaaclab_rl/nerd.py", + ) + for path in paths: + tree = ast.parse(path.read_text()) + for node in ast.walk(tree): + if isinstance(node, (ast.ClassDef, ast.FunctionDef)): + assert node.name not in { + "TransitionCapture", + "HeadingJointStateCodec", + "_write_joint_state", + "NewtonBinding", + "TrajectoryDatasetRecorder", + "_fp32_parallel_nerd_attention", + "AnymalCNeRDEffortExplorationEnvCfg", + "AnymalCNeRDEffortLearnedEnvCfg", + }, path + if isinstance(node, ast.ClassDef) and node.name == "NeRDRecorderCfg": + fields = {item.target.id for item in node.body if isinstance(item, ast.AnnAssign)} + assert fields.isdisjoint({"control", "asset_name"}), "Capture has one applied-effort contract." + if isinstance(node, ast.ClassDef) and node.name == "JointDynamics": + init = next(item for item in node.body if isinstance(item, ast.FunctionDef) and item.name == "__init__") + assert "control" not in {arg.arg for arg in (*init.args.args, *init.args.kwonlyargs)} + imports = [] + if isinstance(node, ast.Import): + imports = [item.name for item in node.names] + elif isinstance(node, ast.ImportFrom) and node.level == 0: + imports = [node.module or ""] + forbidden = ("isaaclab_physicsnemo", "supermimic", "argparse", "gymnasium", "wandb") + if path.is_relative_to(newton): + forbidden += ("isaaclab_rl", "isaaclab_tasks") + if path.is_relative_to(task_core) and path.name != "__init__.py": + forbidden += ("isaaclab_rl", "physicsnemo") + # Registration files legitimately import Gym, but never learner implementations. + if path.name != "__init__.py": + assert not any(name.startswith(forbidden) for name in imports), path + if isinstance(node, ast.Assign): + for target in node.targets: + assert not ( + isinstance(target, ast.Attribute) + and isinstance(target.value, ast.Name) + and target.value.id == "NewtonManager" + ), "Solver state belongs to the simulation's manager instance." + + +def never(env): + return torch.zeros(env.num_envs, dtype=torch.bool, device=env.device) + + +@pytest.mark.parametrize( + "task,network", + ( + ("Isaac-Cartpole-NeRD", "parallel"), + ("Isaac-Velocity-Flat-Anymal-C-NeRD", "parallel"), + ("Isaac-Velocity-Flat-Anymal-C-NeRD", "vector"), + ("Isaac-Velocity-Rough-Anymal-C-NeRD", "parallel"), + ), +) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +def test_record_fit_resume_and_graph_resets(task, network, device, tmp_path): + pytest.importorskip("physicsnemo_robotics") + from isaaclab_newton.physics.nerd import JointDynamics + + from isaaclab_rl.nerd import NeRDTrainCfg, OfflineNeRD + + torch.backends.cuda.matmul.allow_tf32 = False + cfg, _ = resolve_task_config(task + "-Explore", None, overrides=[]) + cfg.scene.num_envs, cfg.sim.device, cfg.seed = 4, device, 7 + if "Rough" in task: + from isaaclab.envs.mdp import reset_root_state_uniform + from isaaclab.terrains import HfRandomUniformTerrainCfg + + cfg.scene.terrain.terrain_generator.num_rows = cfg.scene.terrain.terrain_generator.num_cols = 1 + cfg.scene.terrain.terrain_generator.sub_terrains = { + "rough": HfRandomUniformTerrainCfg(noise_range=(-0.02, 0.02), noise_step=0.005) + } + cfg.events.reset_base.func = reset_root_state_uniform + cfg.events.reset_base.params["pose_range"].update(x=(0.0, 0.0), y=(0.0, 0.0), z=(-0.35, -0.35)) + cfg.recorders.transition.record_to_episodes = True + cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_ALL + cfg.recorders.dataset_export_dir_path = str(tmp_path) + cfg.recorders.export_in_close = True + cfg.recorders.buffer_initial_capacity = 2 + cfg.episode_length_s = 6 * cfg.sim.dt * cfg.decimation + with launch_simulation(cfg): + env = gym.make(task + "-Explore", cfg=cfg).unwrapped + try: + env.reset() + manager = env.sim.backend_manager + capture = env.recorder_manager.get_term("transition") + descriptor = capture.dynamics.descriptor + if descriptor["heading"]: + with pytest.raises(ValueError, match="codec signature"): + JointDynamics(dict(descriptor, heading=False), manager.get_model()) + if "Cartpole" in task: + with pytest.raises(ValueError, match="version 2"): + JointDynamics(dict(descriptor, version=1)) + with pytest.raises(ValueError, match="applied joint efforts"): + JointDynamics(dict(descriptor, control="joint_target_q")) + for gains in (manager.get_model().joint_target_ke, manager.get_model().joint_target_kd): + saved_gains = wp.clone(gains) + gains.fill_(100.0) + with pytest.raises(ValueError, match="solver-side drives"): + JointDynamics(model=manager.get_model(), frame_dt=env.step_dt) + wp.copy(gains, saved_gains) + + def clear_consumed_effort(): + manager.get_control().joint_f.zero_() + + # Capture must retain the solver's input even when a later hook reuses its buffer. + manager.register_post_step_callback(clear_consumed_effort) + recorded_descriptor = dict(descriptor) + if network == "vector": + # A real pre-v3 recording has a layout instead of a portable codec signature. + from dataclasses import fields + + recorded_descriptor.pop("codec") + recorded_descriptor["layout"] = { + field.name: value.tolist() if isinstance(value, torch.Tensor) else value + for field in fields(capture.dynamics.codec.layout) + if field.init and field.name not in ("q_indices", "qd_indices", "semantic_signature") + for value in (getattr(capture.dynamics.codec.layout, field.name),) + } + recorded_descriptor["layout"]["world_count"] = 1 + metadata = {"dynamics": recorded_descriptor, "task": task + "-Explore"} + (tmp_path / "dataset.json").write_text(json.dumps(metadata)) + # Stagger autoresets so prefixes and retained fields exercise different world lengths. + env.episode_length_buf[:] = torch.arange(4, device=device) + seen_terminal = False + seen_contacts = False + with torch.inference_mode(): + for step in range(16): + raw = manager.get_state_0() + before = torch.cat( + (wp.to_torch(raw.joint_q).reshape(4, -1), wp.to_torch(raw.joint_qd).reshape(4, -1)), -1 + ).clone() + actions = torch.full((4, env.action_manager.total_action_dim), 0.1 * (step % 3 - 1), device=device) + if "Cartpole" in task: + # Velocity commands belong to the actuator, and capture must include damping and clipping. + env.scene["robot"].actuators.target_command.set_velocity_index( + value=torch.full((4, 2), 0.25, device=device) + ) + if step == 5: + actions.fill_(10.0) + env.step(actions) + if "Rough" in task: + contacts = capture.latest["contacts"] + valid = contacts[..., 0] == 1 + seen_contacts |= bool(valid.any()) + assert ( + (contacts[..., 1][valid] >= 0) + & (contacts[..., 1][valid] < len(descriptor["contact_entities"])) + ).all() + torch.testing.assert_close( + contacts[..., 7:10][valid].norm(dim=-1), torch.ones_like(contacts[..., 0][valid]) + ) + torch.testing.assert_close(capture.latest["state"], before, rtol=0, atol=0) + if "Cartpole" in task: + expected = (actions[:, 0] * 100 + 10 * (0.25 - before[:, 2])).clamp(-400, 400) + torch.testing.assert_close(capture.latest["control"][:, 0], expected, rtol=1e-6, atol=1e-6) + else: + robot = env.scene["robot"] + names = [ + name.rsplit("/", 1)[-1] + for name in manager.get_model().joint_label[: manager.get_model().joint_count // 4] + ] + names = [name for name in names if name in robot.joint_names] + expected = robot.actuators.applied_effort.torch[ + :, [robot.joint_names.index(name) for name in names] + ] + live = ~env.reset_buf + torch.testing.assert_close(capture.latest["control"][live, 6:], expected[live], rtol=0, atol=0) + if env.reset_buf.any(): + after = torch.cat( + (wp.to_torch(raw.joint_q).reshape(4, -1), wp.to_torch(raw.joint_qd).reshape(4, -1)), -1 + ) + assert not torch.equal(capture.latest["next_state"][env.reset_buf], after[env.reset_buf]) + seen_terminal = True + assert seen_terminal and manager._graph is not None + if "Rough" in task: + assert seen_contacts + # Last continuing sample agrees with cold episode materialization, byte for byte. + for world in range(4): + episode = env.recorder_manager.get_episode(world) + if not episode.is_empty(): + for name, tensor in capture.latest.items(): + assert torch.equal(episode.data["transition"][name][-1], tensor[world]) + env.reset() + assert (capture.latest["terminated"] | capture.latest["truncated"]).all() + env.recorder_manager.export_episodes() + finally: + env.close() + dataset = tmp_path / "dataset.hdf5" + with h5py.File(dataset) as data: + assert len(data["data"]) >= 4 + assert sum(len(episode["transition/state"]) for episode in data["data"].values()) == 64 + + train_cfg = NeRDTrainCfg( + max_updates=3, + batch_size=8, + compile=False, + network={ + "num_heads": 1, + "head_dim": 16, + "intermediate_size": 32, + "head_hidden": 16, + "temporal_attention_spans": (2, 2), + "linear_probe": True, + }, + ) + if "Rough" in task: + train_cfg.network.update(contact_slot_count=2, contact_slot_dim=8) + train_cfg.precision = "bfloat16" + elif network == "vector": + train_cfg.network_type = "physicsnemo_robotics.models.nerd:NewtonNeRDVectorTransformer" + train_cfg.network = { + "num_heads": 2, + "hidden_size": 32, + "head_hidden": 16, + "temporal_depth": 2, + "temporal_attention_spans": (2, 2), + "position_encoding": "relative_lag", + } + elif "Anymal" in task: + train_cfg.precision = "float16" + learner = OfflineNeRD([dataset], train_cfg, device) + assert all(torch.isfinite(torch.tensor(v)) for v in learner.update().values()) + checkpoint = tmp_path / "model.pt" + torch.save(learner.checkpoint(), checkpoint) + resumed = OfflineNeRD( + [dataset], train_cfg, device, checkpoint=torch.load(checkpoint, map_location="cpu", weights_only=True) + ) + assert torch.equal(learner.generator.get_state(), resumed.generator.get_state()) + assert learner.normalizer.fingerprint == resumed.normalizer.fingerprint + assert learner.scaler.state_dict() == resumed.scaler.state_dict() + learner.update() + resumed.update() + for name, tensor in learner.model.state_dict().items(): + # GPU/JIT reduction ordering can change final mantissa bits; resume preserves the update numerically. + torch.testing.assert_close(tensor, resumed.model.state_dict()[name], rtol=1.0e-6, atol=1.0e-8) + assert all(torch.isfinite(torch.tensor(v)) for v in learner.validate().values()) + torch.save(learner.checkpoint(), checkpoint) + + if "Cartpole" in task: + # A two-rank CPU run must match one optimizer step over the combined samples. + # This checks gradient averaging independently of DDP and exercises the actual CLI/save root. + options = tmp_path / "training.json" + options.write_text(json.dumps(to_dict(train_cfg))) + output = tmp_path / "distributed" + script = Path(ISAACLAB_ROOT) / "scripts/learning/nerd/train.py" + command = [sys.executable, "-m", "torch.distributed.run", "--standalone", "--nproc_per_node=2"] + command += [str(script), "--dataset", str(dataset), "--output", str(output)] + command += ["--config", str(options), "--device", "cpu", "--log_interval", "1"] + with (tmp_path / "distributed.log").open("w") as log: + subprocess.run(command + ["--max_updates", "1"], stdout=log, stderr=log, check=True, timeout=120) + distributed = torch.load(output / "last.pt", map_location="cpu", weights_only=True) + oracle = OfflineNeRD([dataset], train_cfg, "cpu") + samples = [] + for rank in range(2): + oracle.generator.manual_seed(train_cfg.seed + rank) + samples.append(oracle.sample(oracle.train_ids, train_cfg.batch_size // 2)) + inputs, targets = (torch.cat([sample[i] for sample in samples]) for i in range(2)) + (oracle.model.predict_last(inputs) - targets).square().mean().backward() + torch.nn.utils.clip_grad_norm_(oracle.model.parameters(), train_cfg.gradient_clip_norm) + oracle.optimizer.step() + for name, tensor in oracle.model.state_dict().items(): + torch.testing.assert_close(tensor, distributed["model"][name], rtol=1e-5, atol=1e-7) + assert len(distributed["generators"]) == 2 + assert not torch.equal(*distributed["generators"]) + with (tmp_path / "distributed.log").open("a") as log: + subprocess.run( + command + ["--max_updates", "2", "--checkpoint", str(output / "last.pt")], + stdout=log, + stderr=log, + check=True, + timeout=120, + ) + continued = torch.load(output / "last.pt", map_location="cpu", weights_only=True) + assert continued["updates"] == 2 + assert all(not torch.equal(a, b) for a, b in zip(distributed["generators"], continued["generators"])) + assert len((output / "metrics.jsonl").read_text().splitlines()) == 2 + + _check_graph_resets(task, device, checkpoint) + + +def _check_graph_resets(task, device, checkpoint): + from isaaclab_newton.physics.nerd_manager import SolverNeRD + + cfg, _ = resolve_task_config(task + "-Learned", None, overrides=[]) + cfg.play_mode() + cfg.scene.num_envs, cfg.sim.device, cfg.seed = 8, device, 7 + if "Rough" in task: + # Unseen terrain and different world count must preserve the robot/material contract. + cfg.scene.terrain.terrain_generator.seed = 99 + cfg.scene.terrain.terrain_generator.num_rows = cfg.scene.terrain.terrain_generator.num_cols = 1 + # Playback imports visual-only geometry absent from headless collection; it is not physics. + cfg.sim.physics.load_visual_shapes = True + cfg.sim.physics.solver_cfg.checkpoint = str(checkpoint) + cfg.sim.physics.solver_cfg.compile = "Anymal" in task + for name, term in vars(cfg.terminations).items(): + if term is not None: + term.func, term.params = never, {} + with launch_simulation(cfg): + env = gym.make(task + "-Learned", cfg=cfg).unwrapped + try: + env.reset() + manager = env.sim.backend_manager + reference = SolverNeRD(manager.get_model(), checkpoint, compile=cfg.sim.physics.solver_cfg.compile) + collision = newton.CollisionPipeline(manager.get_model()) if "Rough" in task else None + contacts = collision.contacts() if collision else None + state = manager.get_model().state() + state.assign(manager.get_state_0()) + action = torch.zeros(8, env.action_manager.total_action_dim, device=device) + if cfg.decimation > 1: + # Save every applied effort, then independently replay four + # single solver steps. ANYdrive's recurrent state stays owned by the real actuator. + efforts = torch.empty(cfg.decimation, manager.get_model().joint_dof_count, device=device) + cursor = torch.zeros(1, dtype=torch.long, device=device) + replay_control = manager.get_model().control() + + def retain_effort(): + with torch.cuda.stream(wp.stream_to_torch(wp.get_stream(device))): + efforts.index_copy_( + 0, cursor.remainder(cfg.decimation), wp.to_torch(manager.get_control().joint_f)[None] + ) + cursor.add_(1) + + manager.register_post_actuator_callback(retain_effort) + addresses = ( + manager._solver.session.buffers.state_history.data_ptr(), + manager._solver.session.buffers.reset_mask.data_ptr(), + ) + with torch.inference_mode(): + for step in range(21): + if step in (4, 9, 13): + ids = torch.tensor([step % 4], device=device) + env.reset(ids) + state.assign(manager.get_state_0()) + mask = wp.zeros(9, dtype=wp.bool, device=device) + wp.to_torch(mask)[ids] = True + reference.reset(state, world_mask=mask) + if step < 2 or step % 2 or not manager.handles_decimation(): + env.step(action) + else: + # Direct graph replay must advance history without executing Python cursor updates. + wp.capture_launch(manager._graph) + if cfg.decimation > 1: + for substep in range(cfg.decimation): + wp.copy(replay_control.joint_f, wp.from_torch(efforts[substep])) + if collision: + collision.collide(state, contacts) + reference.step(state, state, replay_control, contacts, cfg.sim.dt) + else: + reference.step(state, state, manager.get_control(), None, env.step_dt) + wp.synchronize_device(device) + torch.testing.assert_close( + wp.to_torch(manager.get_state_0().joint_q), + wp.to_torch(state.joint_q), + rtol=2e-5, + atol=2e-6, + ) + torch.testing.assert_close( + wp.to_torch(manager.get_state_0().joint_qd), + wp.to_torch(state.joint_qd), + rtol=2e-5, + atol=2e-6, + ) + if step in (0, 4, 13, 20): + solver = manager._solver + operator, buffers = solver.operator, solver.session.buffers + inputs = operator.adapter.encode_inference_history( + buffers.state_history, external_features=buffers.feature_history + ) + inputs = operator.normalizer.normalize("input", inputs) + tokens = None + if "Rough" in task: + from physicsnemo_robotics.dynamics import normalize_contact_tokens + + normalized = normalize_contact_tokens( + operator.normalizer, + "contact", + buffers.contact_history, + operator.adapter.contact_features.signature, + ) + tokens = operator.adapter.model_contacts(operator.model, normalized) + prediction = ( + operator.model.predict_last(inputs) + if tokens is None + else operator.model.predict_last(inputs, contact_tokens=tokens) + ) + delta = operator.normalizer.denormalize("target", prediction) + expected = operator.codec.delta_to_state(buffers.state_history[:, -1], delta) + torch.testing.assert_close(buffers.output, expected, rtol=2e-5, atol=2e-5) + assert manager._graph is not None + assert addresses == ( + manager._solver.session.buffers.state_history.data_ptr(), + manager._solver.session.buffers.reset_mask.data_ptr(), + ) + with pytest.raises(ValueError, match="frame"): + reference.step(state, state, manager.get_control(), None, env.step_dt * 2) + # A checkpoint must reject altered physics even when tensor shapes still agree. + masses = wp.to_torch(manager.get_model().body_mass) + saved_masses = masses.clone() + masses.mul_(1.1) + with pytest.raises(ValueError, match="invalidate"): + reference.notify_model_changed(1) + masses.copy_(saved_masses) + if "Rough" in task: + materials = wp.to_torch(manager.get_model().shape_material_mu) + saved_materials = materials.clone() + materials.mul_(0.8) + with pytest.raises(ValueError, match="invalidate"): + reference.notify_model_changed(1) + materials.copy_(saved_materials) + if "Cartpole" in task: + anchors = wp.to_torch(manager.get_model().joint_X_p) + saved_anchors = anchors.clone() + roots = wp.to_torch(manager.get_model().joint_parent) == -1 + anchors[roots, 3:] = anchors.new_tensor([0.0, 2**-0.5, 0.0, 2**-0.5]) + with pytest.raises(ValueError, match="invalidate"): + reference.notify_model_changed(1) + anchors.copy_(saved_anchors) + flags = wp.to_torch(manager.get_model().body_flags) + saved_flags = flags.clone() + flags.zero_() + with pytest.raises(ValueError, match="invalidate"): + reference.notify_model_changed(1) + flags.copy_(saved_flags) + model = manager.get_model() + pairs = model.shape_collision_filter_pairs_array() + assert len(pairs) + model._set_shape_collision_filter_pairs(()) + with pytest.raises(ValueError, match="invalidate"): + reference.notify_model_changed(1) + model._set_shape_collision_filter_pairs(pairs) + finally: + env.close() diff --git a/source/isaaclab_rl/changelog.d/native-nerd.added.rst b/source/isaaclab_rl/changelog.d/native-nerd.added.rst new file mode 100644 index 000000000000..ed193a8f4f81 --- /dev/null +++ b/source/isaaclab_rl/changelog.d/native-nerd.added.rst @@ -0,0 +1,8 @@ +* Added offline NeRD fitting from native HDF5 episodes, with bounded history sampling, held-out + episode validation, atomic resumable checkpoints and optional synchronous W&B model uploads. +* Supported ParallelNeRD/VectorNeRD, mixed precision with resumable gradient scaling and distributed + fitting with global batches, partitioned validation and per-rank sampler checkpoints. +* Added offline NeRD contact-set fitting with valid-geometry normalization and preserved categorical identities. + +* Shared PhysicsNeMo field normalization between fitting and deployment. Version-3 checkpoints require + refitting existing raw recordings instead of silently reinterpreting earlier heading-normalized weights. diff --git a/source/isaaclab_rl/changelog.d/native-nerd.minor b/source/isaaclab_rl/changelog.d/native-nerd.minor new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_rl/changelog.d/supermimic.added.rst b/source/isaaclab_rl/changelog.d/supermimic.added.rst new file mode 100644 index 000000000000..cd5588a19101 --- /dev/null +++ b/source/isaaclab_rl/changelog.d/supermimic.added.rst @@ -0,0 +1,5 @@ +* Added native SuperMimic training and playback through the unified reinforcement-learning entrypoints, + with MLP and optional one-frame NeRD dynamics, learner-owned replay, resumable checkpoints and W&B uploads. + +* Reused PhysicsNeMo Robotics online normalization for both model architectures, with persistent + device statistics, compatible checkpoint buffers and differentiable imagined policy rollouts. diff --git a/source/isaaclab_rl/changelog.d/supermimic.minor b/source/isaaclab_rl/changelog.d/supermimic.minor new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/api.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/api.py index 8db3c29ed195..2652462473e9 100644 --- a/source/isaaclab_rl/isaaclab_rl/entrypoints/api.py +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/api.py @@ -12,7 +12,7 @@ from .dispatch import run_play_cli, run_random_agent_cli, run_train_cli, run_zero_agent_cli -BackendName = Literal["rl_games", "rlinf", "rsl_rl", "sb3", "skrl", "torchrl"] +BackendName = Literal["rl_games", "rlinf", "rsl_rl", "sb3", "skrl", "supermimic", "torchrl"] """Backends of the unified entrypoints.""" MULTI_GPU_BACKENDS: tuple[BackendName, ...] = ("rl_games", "rsl_rl", "skrl") diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_supermimic.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_supermimic.py new file mode 100644 index 000000000000..3b955bcaf470 --- /dev/null +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_supermimic.py @@ -0,0 +1,150 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""Native SuperMimic policy playback and complete motion-set evaluation.""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import torch + +from isaaclab.app import add_launcher_args, launch_simulation +from isaaclab.utils import update_from_dict + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.contrib.g1_tracking.g1_tracking_env_cfg import TransitionRecorderCfg +from isaaclab_tasks.contrib.g1_tracking.tracking import MAX_REWARD_RATE, OBS_DIM +from isaaclab_tasks.utils import resolve_task_config, setup_preset_cli + +from ...supermimic.networks import Policy +from ..common import ( + add_common_play_args, + apply_env_overrides, + apply_video_recording, + create_isaaclab_env, + pre_launch_video_config, + scoped_torch_backend_flags, +) + + +@torch.no_grad() +def evaluate(env, policy: Policy, real_time: bool = False) -> dict: + """Score each clip once, retaining failed tails in the full-reference-length denominator.""" + native_env = env.unwrapped + command = native_env.command_manager.get_term("motion") + capture = native_env.recorder_manager.get_term("transition") + motion = command.motion + clip_count = len(motion.names) + returns = torch.zeros(clip_count, device=native_env.device) + steps = torch.zeros(clip_count, dtype=torch.long, device=native_env.device) + failed = torch.zeros(clip_count, dtype=torch.bool, device=native_env.device) + for offset in range(0, clip_count, native_env.num_envs): + count = min(native_env.num_envs, clip_count - offset) + command.cfg.evaluation_clip_offset = offset + observations, _ = env.reset() + finished = torch.arange(native_env.num_envs, device=native_env.device) >= count + while not bool(finished.all()): + started = time.perf_counter() + action = policy(observations["policy"], observations["goal"]) + observations, _, _, _, _ = env.step(action) + transition = capture.latest + alive = ~finished + returns[offset : offset + count] += torch.where(alive[:count], transition["r"][:count], 0.0) + steps[offset : offset + count] += alive[:count] + failed[offset : offset + count] |= alive[:count] & transition["term"][:count] + finished |= transition["done"] + if real_time: + time.sleep(max(0.0, native_env.step_dt - (time.perf_counter() - started))) + maximum_step = MAX_REWARD_RATE * native_env.step_dt * native_env.reward_manager.get_term_cfg("tracking").weight + denominators = (motion.clip_len - 1).float() * maximum_step + normalized = returns / denominators + return { + "norm_return": normalized.mean().item(), + "norm_return_framew": (returns.sum() / denominators.sum()).item(), + "n_failed": int(failed.sum()), + "failed": [motion.names[i] for i in failed.nonzero().flatten().tolist()], + "per_clip": dict(zip(motion.names, normalized.tolist(), strict=True)), + "per_clip_steps": dict(zip(motion.names, steps.tolist(), strict=True)), + } + + +def run(argv: list[str]) -> None: + """Restore a native checkpoint's task settings and evaluate its policy with physical simulation.""" + parser = argparse.ArgumentParser(description=__doc__) + add_common_play_args( + parser, agent_default="supermimic_cfg_entry_point", agent_help="SuperMimic agent configuration." + ) + parser.add_argument("--checkpoint", required=True, help="Full or policy-only native SuperMimic checkpoint.") + parser.add_argument("--output", help="Evaluation JSON; defaults to _eval.json.") + parser.add_argument("--motion_dir", help="Evaluate a different directory of reference motions.") + parser.add_argument("--asset_path", help="Relocated USD asset path.") + parser.add_argument("--kinematics_path", help="Relocated matching MJCF asset path.") + add_launcher_args(parser) + args, overrides = setup_preset_cli(parser, argv) + if overrides or args.train_env_cfg or args.frontend != "torch": + parser.error( + "Evaluation restores the saved task and uses its play mode; use explicit relocation/device/seed options." + ) + checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=True) + task = args.task or checkpoint["task"] + if task != checkpoint["task"]: + parser.error("The selected task must match the checkpoint.") + env_cfg, _ = resolve_task_config(task, None, overrides=[]) + env_cfg.recorders.transition = TransitionRecorderCfg() + update_from_dict(env_cfg, checkpoint["env_cfg"]) + apply_env_overrides(args, env_cfg) + if args.motion_dir: + env_cfg.commands.motion.motion_dir = args.motion_dir + if args.asset_path: + env_cfg.scene.robot.spawn.usd_path = args.asset_path + if args.kinematics_path: + env_cfg.commands.motion.kinematics_path = args.kinematics_path + if args.seed is not None: + env_cfg.seed = args.seed + count = sum(path.suffix in (".csv", ".npy") for path in Path(env_cfg.commands.motion.motion_dir).rglob("*")) + if count < 1 or env_cfg.scene.num_envs < 1: + parser.error("Evaluation needs reference motions and a positive world count.") + env_cfg.scene.num_envs = min(env_cfg.scene.num_envs, count) + env_cfg.play_mode() + output = Path(args.output or Path(args.checkpoint).with_suffix("").as_posix() + "_eval.json") + output.parent.mkdir(parents=True, exist_ok=True) + pre_launch_video_config(env_cfg, args) + apply_video_recording(env_cfg, str(output.parent), args, subdir="play", checkpoint_path=args.checkpoint) + torch.manual_seed(env_cfg.seed) + with ( + scoped_torch_backend_flags( + cuda_matmul_allow_tf32=False, cudnn_allow_tf32=True, cudnn_deterministic=False, cudnn_benchmark=False + ), + launch_simulation(env_cfg, args), + ): + env = create_isaaclab_env(task, env_cfg, args, convert_marl_to_single_agent=False) + try: + dtype = torch.bfloat16 if checkpoint["agent_cfg"]["amp"] == "bf16" else None + policy = Policy( + OBS_DIM, + 73 * env_cfg.commands.motion.goal_frames, + hist=env_cfg.observations.policy.history_length, + amp_dtype=dtype, + ).to(env.unwrapped.device) + policy.load_state_dict(checkpoint["pi"]) + policy.eval() + result = evaluate(env, policy, args.real_time) + result.update( + checkpoint=str(Path(args.checkpoint).resolve()), iteration=checkpoint["iteration"], seed=env_cfg.seed + ) + output.write_text(json.dumps(result, indent=2) + "\n") + print( + f"Completed {len(result['per_clip'])} clips: {result['n_failed']} failures, " + f"normalized return {result['norm_return']:.6f}; {output}" + ) + finally: + env.close() diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/train_supermimic.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/train_supermimic.py new file mode 100644 index 000000000000..a4472e78008d --- /dev/null +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/train_supermimic.py @@ -0,0 +1,451 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""SuperMimic training backend of the unified reinforcement learning entrypoint.""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import signal +import time +from contextlib import nullcontext +from datetime import datetime +from pathlib import Path + +import torch + +from isaaclab.app import add_launcher_args, launch_simulation +from isaaclab.utils import string_to_callable, to_dict +from isaaclab.utils.noise import NoiseCfg + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.contrib.g1_tracking.g1_tracking_env_cfg import ( + G1TrackingEnvCfg, + ObservationsCfg, + TerminationsCfg, + TransitionRecorderCfg, +) +from isaaclab_tasks.contrib.g1_tracking.tracking import MAX_REWARD_RATE, OBS_DIM +from isaaclab_tasks.utils import resolve_task_config, setup_preset_cli + +from ...supermimic.learner import Learner +from ...supermimic.networks import WM_IN, WM_OUT, WorldModel +from ..common import ( + add_common_train_args, + apply_env_overrides, + apply_video_recording, + create_isaaclab_env, + dump_train_configs, + enable_cameras_for_video, + pre_launch_video_config, + resolve_seed, + scoped_torch_backend_flags, + wrap_sensor_capture, + write_run_manifest, +) + + +def run(argv: list[str]) -> None: + """Fit dynamics from real transitions and improve a policy through imagined tracking rollouts.""" + parser = argparse.ArgumentParser(description=__doc__) + add_common_train_args( + parser, + agent_default="supermimic_cfg_entry_point", + agent_help="SuperMimic agent configuration.", + include_distributed=False, + ) + parser.add_argument("--checkpoint", help="Full native checkpoint to resume; physical state and replay restart.") + parser.add_argument("--log_dir", help="Output directory; defaults to logs/supermimic//.") + add_launcher_args(parser) + args, overrides = setup_preset_cli(parser, argv) + if args.frontend != "torch": + parser.error("SuperMimic currently supports the native Torch environment frontend.") + env_cfg, agent_cfg = resolve_task_config(args.task, args.agent, overrides=overrides) + _configure_training(args, env_cfg, agent_cfg) + + output = Path( + args.log_dir + or Path("logs/supermimic") / agent_cfg.experiment_name / datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + ) + output.mkdir(parents=True, exist_ok=True) + resumed = torch.load(args.checkpoint, map_location="cpu", weights_only=True) if args.checkpoint else None + # Capture semantic configuration before adding run-specific log paths and video outputs. + saved_env_cfg, saved_agent_cfg = to_dict(env_cfg), to_dict(agent_cfg) + if resumed: + if resumed["task"] != args.task or resumed["env_cfg"] != saved_env_cfg: + raise ValueError("Resume requires the checkpoint's task and environment configuration.") + mutable = { + "max_iterations", + "max_hours", + "save_interval", + "save_interval_seconds", + "log_interval", + "wandb_project", + "wandb_entity", + "wandb_run_id", + "wandb_name", + } + changed = { + key for key, value in saved_agent_cfg.items() if key not in mutable and resumed["agent_cfg"][key] != value + } + if changed or resumed["iteration"] >= agent_cfg.max_iterations: + raise ValueError(f"Resume requires unchanged learning settings and a later iteration limit: {changed}.") + identity = resumed.get("wandb") if resumed else None + if identity and agent_cfg.wandb_run_id is not None and agent_cfg.wandb_run_id != identity["run_id"]: + identity = None + if identity: + for key in ("project", "entity", "run_id"): + value = getattr(agent_cfg, f"wandb_{key}") + if value is not None and value != identity[key]: + raise ValueError(f"Resuming a W&B run must preserve its {key}.") + setattr(agent_cfg, f"wandb_{key}", identity[key]) + logging = nullcontext(None) + if agent_cfg.wandb_project: + import wandb + + logging = wandb.init( + project=agent_cfg.wandb_project, + entity=agent_cfg.wandb_entity, + id=agent_cfg.wandb_run_id, + name=agent_cfg.wandb_name, + resume="allow", + dir=str(output), + config={"task": args.task, "agent": saved_agent_cfg, "env": saved_env_cfg}, + allow_val_change=True, + ) + env_cfg.log_dir = str(output) + enable_cameras_for_video(args) + pre_launch_video_config(env_cfg, args) + apply_video_recording(env_cfg, str(output), args) + write_run_manifest(str(output), library="supermimic", task=args.task, metadata={"agent": args.agent}) + dump_train_configs(str(output), env_cfg, agent_cfg) + torch.manual_seed(agent_cfg.seed) + with ( + logging as run, + scoped_torch_backend_flags( + cuda_matmul_allow_tf32=False, cudnn_allow_tf32=True, cudnn_deterministic=False, cudnn_benchmark=False + ), + launch_simulation(env_cfg, args), + ): + if run: + run.define_metric("iteration") + run.define_metric("*", step_metric="iteration") + identity = {"entity": run.entity, "project": run.project, "run_id": run.id, "url": run.url} + (output / "wandb-run.json").write_text(json.dumps(identity, indent=2) + "\n") + env = create_isaaclab_env(args.task, env_cfg, args, convert_marl_to_single_agent=False) + env = wrap_sensor_capture(env, str(output), args) + metadata = { + "task": args.task, + "agent_cfg": saved_agent_cfg, + "env_cfg": saved_env_cfg, + "wandb": identity, + "state_layout": "root-link-xyzw-v-w-q-qd", + } + try: + _train(env, agent_cfg, output, resumed, run, metadata) + finally: + env.close() + + +def _configure_training(args, env_cfg, agent_cfg): + """Apply native overrides and validate that real and imagined task semantics agree.""" + if not isinstance(env_cfg, G1TrackingEnvCfg): + raise ValueError("SuperMimic currently requires the G1 motion-tracking task's analytical objective.") + apply_env_overrides(args, env_cfg) + if args.seed is not None: + agent_cfg.seed = resolve_seed(args.seed) + if args.max_iterations is not None: + agent_cfg.max_iterations = args.max_iterations + env_cfg.seed = agent_cfg.seed + if env_cfg.recorders.transition is None: + env_cfg.recorders.transition = TransitionRecorderCfg() + if min(agent_cfg.buffer_steps, agent_cfg.warmup_steps) < agent_cfg.world_model_window: + raise ValueError("Replay capacity and warmup must each cover a full world-model window.") + if ( + min( + agent_cfg.max_iterations, + agent_cfg.horizon, + agent_cfg.world_model_window, + agent_cfg.world_model_batches, + agent_cfg.world_model_batch_size, + agent_cfg.policy_batch_size, + agent_cfg.critic_epochs, + agent_cfg.critic_minibatches, + agent_cfg.learning_rate_decay_iterations, + agent_cfg.save_interval, + agent_cfg.log_interval, + env_cfg.observations.policy.history_length, + env_cfg.commands.motion.goal_frames, + ) + < 1 + ): + raise ValueError("Learning counts, observation history and goal lengths must be positive.") + if agent_cfg.critic_minibatches > agent_cfg.horizon * agent_cfg.policy_batch_size: + raise ValueError("Critic minibatches must contain samples.") + if not 0 <= env_cfg.commands.motion.gamma < 1 or not 0 <= agent_cfg.gae_lambda <= 1: + raise ValueError("Discount must be in [0, 1) and GAE lambda in [0, 1].") + if not math.isfinite(agent_cfg.action_noise) or agent_cfg.action_noise < 0: + raise ValueError("Action noise must be finite and nonnegative.") + if agent_cfg.world_model_only_iterations < 0: + raise ValueError("World-model-only iterations must be nonnegative.") + for value in (agent_cfg.update_to_data_ratio, agent_cfg.learning_rate, agent_cfg.save_interval_seconds): + if not math.isfinite(value) or value <= 0: + raise ValueError("Learning rates, update/data ratio and checkpoint intervals must be finite and positive.") + if agent_cfg.max_hours is not None and (not math.isfinite(agent_cfg.max_hours) or agent_cfg.max_hours <= 0): + raise ValueError("The training-hour budget must be finite and positive.") + if agent_cfg.amp not in ("none", "bf16") or agent_cfg.world_model not in ("mlp", "nerd-one-frame"): + raise ValueError("Select amp=none/bf16 and world_model=mlp/nerd-one-frame.") + if env_cfg.sim.dt * env_cfg.decimation != 0.02: + raise ValueError("The reference motions and world model require a 20 ms control period.") + # The same analytical observations are used inside imagined rollouts. Reject changes that + # preserve tensor width but alter their meaning; native history and stateless noise remain configurable. + template = ObservationsCfg() + for name in ("policy", "goal"): + expected, actual = to_dict(getattr(template, name)), to_dict(getattr(env_cfg.observations, name)) + if name == "policy": + for group in (expected, actual): + group.pop("history_length") + group.pop("enable_corruption") + for term in group.values(): + if isinstance(term, dict): + term.pop("noise", None) + if actual != expected or tuple(actual) != tuple(expected): + raise ValueError(f"Analytical {name} observations require the task's declared functions and layout.") + if to_dict(env_cfg.terminations) != to_dict(TerminationsCfg()): + raise ValueError("Analytical termination requires the task's declared failure, clip-end and timeout terms.") + + +def _train(env, agent_cfg, output, resumed, run, metadata): + """Own the learning schedule, checkpoint publication and signal handling for one environment.""" + if run: + import wandb + + previous_handlers = {} + try: + native_env = env.unwrapped + command = native_env.command_manager.get_term("motion") + capture = native_env.recorder_manager.get_term("transition") + observations, _ = env.reset() + from isaaclab_tasks.contrib.g1_tracking.mdp import tracking_reward + + reward_cfg = native_env.reward_manager.get_term_cfg("tracking") + reward_func = string_to_callable(reward_cfg.func) if isinstance(reward_cfg.func, str) else reward_cfg.func + if ( + native_env.reward_manager.active_terms != ["tracking"] + or reward_func is not tracking_reward + or reward_cfg.weight <= 0 + ): + raise ValueError("The imagined objective requires the task's single tracking reward with positive weight.") + noise_terms, offset = [], 0 + group = native_env.cfg.observations.policy + for name, shape in zip( + native_env.observation_manager.active_terms["policy"], + native_env.observation_manager.group_obs_term_dim["policy"], + strict=True, + ): + term = getattr(group, name) + noise = term.noise if group.enable_corruption else None + if noise is not None and not isinstance(noise, NoiseCfg): + raise ValueError("Analytical observations require stateless policy noise.") + fn = ( + None if noise is None else string_to_callable(noise.func) if isinstance(noise.func, str) else noise.func + ) + noise_terms.append((slice(offset, offset + shape[-1]), fn, noise)) + offset += shape[-1] + if offset != OBS_DIM or int(command.motion.clip_len.max()) <= agent_cfg.horizon: + raise ValueError("Observation width or motion lengths are incompatible with the learner.") + network = None + if agent_cfg.world_model == "nerd-one-frame": + from physicsnemo_robotics.models.parallel_nerd import ParallelNeRD + + network = ParallelNeRD(WM_IN, WM_OUT, temporal_attention_spans=(1, 1, 1, 1)) + torch.nn.init.zeros_(network.head_output.weight) + torch.nn.init.zeros_(network.head_output.bias) + dtype = torch.bfloat16 if agent_cfg.amp == "bf16" else None + world_model = WorldModel(network=network, amp_dtype=dtype).to(native_env.device) + learner = Learner( + agent_cfg, + command.mdp, + capture.latest, + tuple(noise_terms), + native_env.step_dt * reward_cfg.weight, + world_model, + goal_frames=command.cfg.goal_frames, + history=group.history_length, + gamma=command.cfg.gamma, + ) + first_iteration, prior_wall_time, sim_steps = 0, 0.0, 0 + if resumed: + for name in ("wm", "pi", "critic", "opt_wm", "opt_pi", "opt_v"): + getattr(learner, name).load_state_dict(resumed[name]) + command.bin_return.copy_(resumed["bin_return"]) + command.bin_weight.copy_(resumed["bin_weight"]) + first_iteration, prior_wall_time, sim_steps = ( + resumed["iteration"], + resumed["wall_time"], + resumed["sim_steps"], + ) + del resumed + stop_signal = 0 + + def request_stop(signum, frame): + nonlocal stop_signal + stop_signal = signum + + previous_handlers = {sig: signal.signal(sig, request_stop) for sig in (signal.SIGTERM, signal.SIGINT)} + stats = {name: [] for name in ("reward", "length", "return", "failed", "done")} + + @torch.no_grad() + def collect(count): + nonlocal observations, sim_steps + for _ in range(count): + observations, _, _, _, extras = env.step(learner.act(observations)) + transition = capture.latest + learner.observe(transition) + done = transition["done"] + done_count = done.sum() + stats["reward"].append(transition["r"].mean() / (MAX_REWARD_RATE * learner.reward_scale)) + stats["length"].append(((transition["ep_step"] + 1) * done).sum()) + stats["return"].append(extras["log"].get("Metrics/motion/return", 0.0) * done_count) + stats["failed"].append(transition["term"].sum()) + stats["done"].append(done_count) + sim_steps += native_env.num_envs + + collect(agent_cfg.warmup_steps) + training_start = log_start = checkpoint_time = time.time() + step_accumulator = 0.0 + aggregate = {name: [] for name in ("wm", "wm_last", "pi", "v", "imagined_reward", "gradient", "survival")} + columns = ( + "iteration", + "wall_h", + "sim_steps", + "wm_loss", + "wm_last_error", + "pi_loss", + "v_loss", + "imagined_reward", + "actor_gradient", + "survival", + "sim_reward", + "episode_length", + "episode_return", + "failure_rate", + "iterations_per_second", + "learning_rate", + ) + log_path = output / "train_log.csv" + if first_iteration and log_path.exists(): + lines = log_path.read_text().splitlines() + lines = lines[:1] + [line for line in lines[1:] if int(float(line.split(",")[0])) <= first_iteration] + log_path.write_text("\n".join(lines) + "\n") + else: + log_path.write_text(",".join(columns) + "\n") + with log_path.open("a") as log: + for iteration in range(first_iteration, agent_cfg.max_iterations): + fraction = min(iteration / agent_cfg.learning_rate_decay_iterations, 1.0) + lr = agent_cfg.learning_rate * (0.03 + 0.97 * 0.5 * (1 + math.cos(math.pi * fraction))) + for optimizer in (learner.opt_wm, learner.opt_pi, learner.opt_v): + for group_cfg in optimizer.param_groups: + group_cfg["lr"] = lr + step_accumulator += 1.0 / agent_cfg.update_to_data_ratio + count = int(step_accumulator) + step_accumulator -= count + collect(count) + for _ in range(agent_cfg.world_model_batches): + loss, error = learner.wm_update() + aggregate["wm"].append(loss) + aggregate["wm_last"].append(error) + if iteration >= agent_cfg.world_model_only_iterations: + results = learner.policy_update() + for name, value in zip( + ("pi", "v", "imagined_reward", "gradient", "survival"), results, strict=True + ): + aggregate[name].append(value) + if (iteration + 1) % agent_cfg.log_interval == 0: + means = [float(torch.stack(values).mean()) if values else 0.0 for values in aggregate.values()] + totals = { + name: float(torch.stack(values).sum()) if values else 0.0 for name, values in stats.items() + } + done_count = max(totals["done"], 1.0) + wall_h = (prior_wall_time + time.time() - training_start) / 3600 + row = [ + iteration + 1, + wall_h, + sim_steps, + *means, + totals["reward"] / max(len(stats["reward"]), 1), + totals["length"] / done_count, + totals["return"] / done_count, + totals["failed"] / done_count, + agent_cfg.log_interval / (time.time() - log_start), + lr, + ] + if not all(math.isfinite(value) for value in row): + raise FloatingPointError(f"Nonfinite training statistics at update {iteration + 1}: {row}") + log.write(",".join(f"{value:.6g}" for value in row) + "\n") + log.flush() + if run: + run.log(dict(zip(columns, row, strict=True))) + print( + f"Iteration {iteration + 1}: dynamics={means[0]:.4f}, policy={means[2]:.4f}, " + f"value={means[3]:.4f}, hours={wall_h:.3f}", + flush=True, + ) + log_start = time.time() + for values in (*stats.values(), *aggregate.values()): + values.clear() + elapsed = prior_wall_time + time.time() - training_start + stop = agent_cfg.max_hours is not None and elapsed >= agent_cfg.max_hours * 3600 + milestone = ( + (iteration + 1) % agent_cfg.save_interval == 0 or iteration + 1 == agent_cfg.max_iterations or stop + ) + if ( + milestone + or stop_signal + or iteration == first_iteration + or time.time() - checkpoint_time >= agent_cfg.save_interval_seconds + ): + checkpoint = { + **metadata, + **{ + name: getattr(learner, name).state_dict() + for name in ("wm", "pi", "critic", "opt_wm", "opt_pi", "opt_v") + }, + "iteration": iteration + 1, + "wall_time": elapsed, + "sim_steps": sim_steps, + "bin_return": command.bin_return, + "bin_weight": command.bin_weight, + "saved_at": datetime.now().isoformat(), + } + torch.save(checkpoint, output / "last.pt.tmp") + os.replace(output / "last.pt.tmp", output / "last.pt") + if milestone: + policy_checkpoint = { + key: value + for key, value in checkpoint.items() + if key not in ("wm", "critic", "opt_wm", "opt_pi", "opt_v") + } + path = output / f"policy_{iteration + 1:06d}.pt" + torch.save(policy_checkpoint, path.with_suffix(".tmp")) + os.replace(path.with_suffix(".tmp"), path) + if run: + artifact = wandb.Artifact( + f"run-{run.id}-checkpoint", type="model", metadata={"iteration": iteration + 1} + ) + artifact.add_file(str(output / "last.pt"), name="last.pt") + run.log_artifact(artifact, aliases=["latest"]).wait(timeout=180) + run.summary["durable_checkpoint_iteration"] = iteration + 1 + checkpoint_time = time.time() + if stop_signal: + raise SystemExit(128 + stop_signal) + if stop: + break + finally: + for sig, handler in previous_handlers.items(): + signal.signal(sig, handler) diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/dispatch.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/dispatch.py index aff2b314193e..b50d2378cf69 100644 --- a/source/isaaclab_rl/isaaclab_rl/entrypoints/dispatch.py +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/dispatch.py @@ -33,6 +33,7 @@ "rsl_rl": "isaaclab_rl.entrypoints.backends.train_rsl_rl", "sb3": "isaaclab_rl.entrypoints.backends.train_sb3", "skrl": "isaaclab_rl.entrypoints.backends.train_skrl", + "supermimic": "isaaclab_rl.entrypoints.backends.train_supermimic", "torchrl": "isaaclab_rl.entrypoints.backends.train_torchrl", }, "play": { @@ -41,6 +42,7 @@ "rsl_rl": "isaaclab_rl.entrypoints.backends.play_rsl_rl", "sb3": "isaaclab_rl.entrypoints.backends.play_sb3", "skrl": "isaaclab_rl.entrypoints.backends.play_skrl", + "supermimic": "isaaclab_rl.entrypoints.backends.play_supermimic", "torchrl": "isaaclab_rl.entrypoints.backends.play_torchrl", }, } diff --git a/source/isaaclab_rl/isaaclab_rl/nerd.py b/source/isaaclab_rl/isaaclab_rl/nerd.py new file mode 100644 index 000000000000..90f131295e87 --- /dev/null +++ b/source/isaaclab_rl/isaaclab_rl/nerd.py @@ -0,0 +1,296 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Offline NeRD fitting from native RecorderManager episodes.""" + +from __future__ import annotations + +import hashlib +import json +import math +from pathlib import Path + +import h5py +import torch +import torch.distributed as dist +from isaaclab_newton.physics.nerd import JointDynamics +from physicsnemo_robotics.dynamics import FieldNormalizer, RigidContactSetEncoder + +from isaaclab.utils import configclass, string_to_callable + + +@configclass +class NeRDTrainCfg: + """Offline model fitting settings; the dataset owns state, control and frame conventions.""" + + max_updates: int = 30000 + batch_size: int = 1024 + learning_rate: float = 1.0e-3 + final_learning_rate: float = 1.0e-4 + learning_rate_decay_updates: int = 30000 + gradient_clip_norm: float = 1.0 + validation_fraction: float = 0.1 + seed: int = 7 + compile: bool = True + network_type: str = "physicsnemo_robotics.models.parallel_nerd:ParallelNeRD" + """Network class with a temporal_attention_spans configuration and a predict_last method.""" + precision: str = "float32" + """Projection precision: float32, float16 or bfloat16; parameters and losses remain FP32.""" + max_gradient_retries: int = 8 + network: dict = { + "num_heads": 4, + "head_dim": 64, + "intermediate_size": 640, + "head_hidden": 256, + "temporal_attention_spans": (4, 3, 3, 3), + "rms_norm_eps": 1.0e-6, + "linear_probe": True, + } + + +class _PredictLast(torch.nn.Module): + """Expose the newest prediction through forward so DDP owns gradient synchronization.""" + + def __init__(self, network): + super().__init__() + self.network = network + + def forward(self, inputs, contacts=None): + return ( + self.network.predict_last(inputs, contacts) if contacts is not None else self.network.predict_last(inputs) + ) + + +class OfflineNeRD: + """Own flattened transition tensors, boundary-aware history sampling and model optimization. + + Episodes, rather than overlapping windows, are split into train and validation sets. The + terminal transition is retained. Every frame after an episode boundary starts a repeated + first-frame prefix, matching deployed solver resets. Storage grows with transitions, not + with the product of transitions and history length. + """ + + def __init__(self, datasets: list[Path], cfg: NeRDTrainCfg, device: str, *, checkpoint=None): + if cfg.max_updates < 1 or cfg.batch_size < 1 or not 0 < cfg.validation_fraction < 1: + raise ValueError("Positive update/batch counts and a validation fraction in (0, 1) are required.") + if cfg.learning_rate_decay_updates < 1 or any( + not math.isfinite(value) or value <= 0 + for value in (cfg.learning_rate, cfg.final_learning_rate, cfg.gradient_clip_norm) + ): + raise ValueError("Learning rates, decay duration and gradient clip norm must be positive and finite.") + self.cfg, self.device = cfg, device + self.rank = dist.get_rank() if dist.is_initialized() else 0 + self.world_size = dist.get_world_size() if dist.is_initialized() else 1 + if cfg.batch_size % self.world_size: + raise ValueError("The global batch size must be divisible by the distributed world size.") + if cfg.precision not in ("float32", "float16", "bfloat16") or cfg.max_gradient_retries < 1: + raise ValueError("Select float32, float16 or bfloat16 and a positive gradient retry limit.") + self.dtype = getattr(torch, cfg.precision) + self.device_type = torch.device(device).type + if self.dtype == torch.float16 and self.device_type != "cuda": + raise ValueError("FP16 fitting requires CUDA.") + self.scaler = torch.amp.GradScaler("cuda", enabled=self.dtype == torch.float16) + self.generator = torch.Generator(device=device).manual_seed(cfg.seed + self.rank) + self.updates = 0 + self.identity = [] + self.dynamics = None + episodes = [] + for path in datasets: + path = Path(path) + metadata = json.loads(path.with_suffix(".json").read_text()) + descriptor = metadata["dynamics"] + if self.dynamics is None: + self.dynamics = JointDynamics(descriptor) + elif self.dynamics.descriptor != descriptor: + raise ValueError("Datasets must use the same dynamics representation and physical model.") + with path.open("rb") as stream: + fingerprint = hashlib.file_digest(stream, "sha256").hexdigest() + if fingerprint in self.identity: + raise ValueError("Duplicate datasets would leak episodes between training and validation.") + self.identity.append(fingerprint) + with h5py.File(path, "r") as file: + for name in file["data"]: + data = file["data"][name]["transition"] + state, successor, control, contacts = ( + torch.from_numpy(data[key][:]).to(device) + for key in ("state", "next_state", "control", "contacts") + ) + if not len(state): + continue + if not all(torch.isfinite(value).all() for value in (state, successor, control, contacts)): + raise ValueError(f"Nonfinite transition in {path}:{name}.") + if not torch.equal(successor[:-1], state[1:]): + raise ValueError( + f"Discontinuous episode in {path}:{name}; a reset cannot enter a training window." + ) + if data["terminated"][:-1].any() or data["truncated"][:-1].any(): + raise ValueError(f"Episode boundary inside {path}:{name}.") + inputs = self.dynamics.inputs(state, control, contacts) + targets = self.dynamics.codec.state_to_delta(state, successor) + episodes.append((inputs, targets, contacts)) + if len(episodes) < 2: + raise ValueError("Record at least two nonempty episodes to split training and validation.") + order = torch.randperm(len(episodes), generator=torch.Generator().manual_seed(cfg.seed)).tolist() + validation_count = max(1, min(len(episodes) - 1, round(cfg.validation_fraction * len(episodes)))) + self.validation_episodes = order[:validation_count] + self.train_episodes = order[validation_count:] + self.inputs = torch.cat([item[0] for item in episodes]) + self.targets = torch.cat([item[1] for item in episodes]) + self.contacts = None + lengths = torch.tensor([len(item[0]) for item in episodes], device=device) + starts = lengths.cumsum(0) - lengths + self.starts = torch.repeat_interleave(starts, lengths) + ids = torch.repeat_interleave(torch.arange(len(episodes), device=device), lengths) + held_out = torch.isin(ids, ids.new_tensor(self.validation_episodes)) + self.train_ids, self.validation_ids = (~held_out).nonzero().flatten(), held_out.nonzero().flatten() + contacts = torch.cat([item[2] for item in episodes]) if self.dynamics.max_contacts else None + if checkpoint is None: + train_x, train_y = self.inputs[self.train_ids], self.targets[self.train_ids] + means, stds = ( + {"input": train_x.mean(0), "target": train_y.mean(0)}, + {"input": train_x.std(0, correction=0), "target": train_y.std(0, correction=0)}, + ) + if contacts is not None: + training = contacts[self.train_ids] + geometry = training[..., 4:][training[..., 0] == 1] + if not len(geometry): + raise ValueError("Contact-conditioned training requires contacts in its training split.") + means["contact"] = torch.cat((geometry.new_zeros(4), geometry.mean(0))) + stds["contact"] = torch.cat((geometry.new_ones(4), geometry.std(0, correction=0))) + self.normalizer = FieldNormalizer(means, stds) + else: + if ( + checkpoint["format"] != "isaaclab-nerd-v3" + or checkpoint["datasets"] != self.identity + or checkpoint["dynamics"] != self.dynamics.descriptor + or checkpoint["train_episodes"] != self.train_episodes + or checkpoint["network_cfg"] != cfg.network + or checkpoint["network_type"] != cfg.network_type + or len(checkpoint["generators"]) != self.world_size + ): + raise ValueError( + "Resume requires v3 canonical inputs and the same datasets, split, network and rank count." + ) + self.normalizer = FieldNormalizer.from_dict(checkpoint["normalizer"]) + self.inputs = self.normalizer.normalize("input", self.inputs) + self.targets = self.normalizer.normalize("target", self.targets) + if contacts is not None: + # IDs/masks are categorical: normalize geometry, then restore them bit-exactly. + normalized = self.normalizer.normalize_with_passthrough("contact", contacts, (0, 1, 2, 3)) + normalized = torch.where(contacts[..., :1] == 1, normalized, 0.0) + self.contacts = RigidContactSetEncoder.to_variable_contact_set(normalized) + history = 1 + sum(span - 1 for span in cfg.network["temporal_attention_spans"]) + self.offsets = torch.arange(1 - history, 1, device=device) + torch.manual_seed(cfg.seed) + self.model = string_to_callable(cfg.network_type)( + self.dynamics.input_dim, + self.dynamics.codec.prediction_shape[0], + **cfg.network, + **self.dynamics.contact_model_cfg, + ).to(device) + self.predict = _PredictLast(self.model) + if self.world_size > 1: + device_ids = [torch.device(device).index] if self.device_type == "cuda" else None + self.predict = torch.nn.parallel.DistributedDataParallel(self.predict, device_ids=device_ids) + if cfg.compile: + self.predict = torch.compile(self.predict) + self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=cfg.learning_rate, weight_decay=0.0) + if checkpoint is not None: + self.model.load_state_dict(checkpoint["model"]) + self.optimizer.load_state_dict(checkpoint["optimizer"]) + self.generator.set_state(checkpoint["generators"][self.rank].cpu()) + self.scaler.load_state_dict(checkpoint["scaler"]) + self.updates = checkpoint["updates"] + + def sample(self, ids, count): + """Draw transitions uniformly, with a reset prefix that never crosses episode boundaries.""" + selected = ids[torch.randint(len(ids), (count,), device=self.device, generator=self.generator)] + history = (selected[:, None] + self.offsets).clamp_min(self.starts[selected, None]) + return ( + self.inputs[history], + self.targets[selected], + self.contacts[history] if self.contacts is not None else None, + ) + + def update(self): + self.model.train() + self.optimizer.param_groups[0]["lr"] = self.cfg.learning_rate + ( + self.cfg.final_learning_rate - self.cfg.learning_rate + ) * min(self.updates / self.cfg.learning_rate_decay_updates, 1.0) + x, y, contacts = self.sample(self.train_ids, self.cfg.batch_size // self.world_size) + for _ in range(self.cfg.max_gradient_retries): + self.optimizer.zero_grad(set_to_none=True) + with torch.autocast(self.device_type, dtype=self.dtype, enabled=self.dtype != torch.float32): + prediction = self.predict(x, contacts) + loss = (prediction.float() - y).square().mean() + finite = torch.isfinite(loss).to(torch.int32) + if self.world_size > 1: + dist.all_reduce(finite, op=dist.ReduceOp.MIN) + if not finite: + raise FloatingPointError(f"Nonfinite NeRD training loss at update {self.updates}.") + self.scaler.scale(loss).backward() + self.scaler.unscale_(self.optimizer) + norm = torch.nn.utils.clip_grad_norm_( + self.model.parameters(), self.cfg.gradient_clip_norm, error_if_nonfinite=not self.scaler.is_enabled() + ) + scale = self.scaler.get_scale() + self.scaler.step(self.optimizer) + self.scaler.update() + if self.scaler.get_scale() >= scale: + break + else: + raise FloatingPointError(f"NeRD gradients overflowed repeatedly at update {self.updates}.") + if self.world_size > 1: + dist.all_reduce(loss.detach()) + loss.detach().div_(self.world_size) + self.updates += 1 + return {"train/normalized_mse": float(loss.detach()), "train/gradient_norm": float(norm)} + + @torch.no_grad() + def validate(self): + self.model.eval() + totals = torch.zeros(3, dtype=torch.float64, device=self.device) + for ids in self.validation_ids[self.rank :: self.world_size].split(self.cfg.batch_size): + if not len(ids): + continue + history = (ids[:, None] + self.offsets).clamp_min(self.starts[ids, None]) + with torch.autocast(self.device_type, dtype=self.dtype, enabled=self.dtype != torch.float32): + contacts = self.contacts[history] if self.contacts is not None else None + prediction = ( + self.model.predict_last(self.inputs[history], contacts) + if contacts is not None + else self.model.predict_last(self.inputs[history]) + ) + error = prediction.float() - self.targets[ids] + totals[0] += error.square().sum() + totals[1] += (error * (self.normalizer.stds["target"].to(error) + self.normalizer.eps)).square().sum() + totals[2] += error.numel() + if self.world_size > 1: + dist.all_reduce(totals) + return { + "validation/normalized_mse": float(totals[0] / totals[2]), + "validation/delta_mse": float(totals[1] / totals[2]), + } + + def checkpoint(self): + """Return a deployable and resumable snapshot; the composition root publishes it atomically.""" + generators = [self.generator.get_state().cpu()] + if self.world_size > 1: + generators = [None] * self.world_size + dist.all_gather_object(generators, self.generator.get_state().cpu()) + return { + "format": "isaaclab-nerd-v3", + "dynamics": self.dynamics.descriptor, + "datasets": self.identity, + "train_episodes": self.train_episodes, + "network_cfg": self.cfg.network, + "network_type": str(self.cfg.network_type), + "model": self.model.state_dict(), + "optimizer": self.optimizer.state_dict(), + "normalizer": self.normalizer.to_dict(), + "updates": self.updates, + "generators": generators, + "scaler": self.scaler.state_dict(), + } diff --git a/source/isaaclab_rl/isaaclab_rl/supermimic/__init__.py b/source/isaaclab_rl/isaaclab_rl/supermimic/__init__.py new file mode 100644 index 000000000000..d0e155aa7944 --- /dev/null +++ b/source/isaaclab_rl/isaaclab_rl/supermimic/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""SuperMimic dynamics fitting and differentiable policy learning.""" diff --git a/source/isaaclab_rl/isaaclab_rl/supermimic/learner.py b/source/isaaclab_rl/isaaclab_rl/supermimic/learner.py new file mode 100644 index 000000000000..4adee693220d --- /dev/null +++ b/source/isaaclab_rl/isaaclab_rl/supermimic/learner.py @@ -0,0 +1,290 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""Replay ownership, supervised dynamics fitting and differentiable actor/critic updates.""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +from isaaclab_tasks.contrib.g1_tracking.tracking import MAX_REWARD_RATE, OBS_DIM, unflatten_state + +from .networks import CRITIC_EXTRA, INV_DIM, Critic, Policy, WorldModel, critic_features, normalize_xi, state_difference +from .supermimic_cfg import SuperMimicCfg + + +class ReplayBuffer: + """Learner-owned time/world ring; sampled model windows cannot cross physical discontinuities.""" + + def __init__(self, capacity: int, transition: dict[str, torch.Tensor]): + if capacity < 1: + raise ValueError("Replay capacity must be positive.") + self.capacity, self.size, self.pointer = (capacity, 0, 0) + self.num_envs, self.device = (transition["s"].shape[0], transition["s"].device) + fields = ("s", "s1", "a", "a_prev", "xi", "f", "ep_step", "done") + self.data = {key: transition[key].new_empty((capacity, *transition[key].shape)) for key in fields} + + @torch.no_grad() + def add(self, transition): + for key, data in self.data.items(): + data[self.pointer].copy_(transition[key]) + self.pointer = (self.pointer + 1) % self.capacity + self.size = min(self.size + 1, self.capacity) + + def _row(self, age): + return (self.pointer - self.size + age) % self.capacity + + def sample_windows(self, batch_size, length): + if self.size < length: + raise ValueError("Replay does not yet contain a full model window.") + start = torch.randint(self.size - length + 1, (batch_size,), device=self.device) + env = torch.randint(self.num_envs, (batch_size, 1), device=self.device) + rows = self._row(start[:, None] + torch.arange(length, device=self.device)) + out = {key: self.data[key][rows, env] for key in ("s", "s1", "a", "xi", "done")} + steps, frames = (self.data["ep_step"][rows, env], self.data["f"][rows, env]) + connected = ( + ~out["done"][:, :-1] + & (steps[:, 1:] == steps[:, :-1] + 1) + & (frames[:, 1:] == frames[:, :-1] + 1) + & (out["s"][:, 1:] == out["s1"][:, :-1]).all(-1) + ) + out["mask"] = torch.cat((torch.ones(batch_size, 1, device=self.device), connected.float()), -1).cumprod(-1) + return out + + def sample_starts(self, batch_size, history_length): + if not self.size: + raise ValueError("Cannot sample an empty replay buffer.") + age = torch.randint(self.size, (batch_size,), device=self.device) + env = torch.randint(self.num_envs, (batch_size,), device=self.device) + rows = self._row(age) + steps = self.data["ep_step"][rows, env].long() + back = torch.arange(history_length - 1, -1, -1, device=self.device) + back = torch.minimum(torch.minimum(back[None], steps[:, None]), age[:, None]) + history_rows = self._row(age[:, None] - back) + fields = ("s", "a_prev", "xi", "f") + start = {key: self.data[key][rows, env] for key in fields} + history = {key: self.data[key][history_rows, env[:, None]] for key in fields} + return (start, history) + + +class Learner: + """Own models, normalization, replay and optimization, independent of simulator stepping.""" + + def __init__( + self, + cfg: SuperMimicCfg, + mdp, + transition, + noise_terms, + reward_scale, + world_model: WorldModel, + *, + goal_frames: int, + history: int, + gamma: float, + ): + self.cfg, self.mdp, self.motion = (cfg, mdp, mdp.motion) + self.gamma = gamma + self.device = transition["s"].device + self.noise_terms, self.reward_scale = (noise_terms, reward_scale) + if not 0 <= gamma < 1: + raise ValueError("Discount must lie in [0, 1).") + self.goal_frames, self.history, self.horizon = (goal_frames, history, cfg.horizon) + dtype = torch.bfloat16 if cfg.amp == "bf16" else None + self.wm = world_model + self.pi = Policy(OBS_DIM, 73 * goal_frames, hist=history, amp_dtype=dtype).to(self.device) + self.critic = Critic(INV_DIM + CRITIC_EXTRA + 73 * goal_frames, amp_dtype=dtype).to(self.device) + self.opt_wm = torch.optim.Adam(self.wm.parameters(), lr=cfg.learning_rate) + self.opt_pi = torch.optim.Adam(self.pi.parameters(), lr=cfg.learning_rate, betas=(0.9, 0.95)) + self.opt_v = torch.optim.Adam(self.critic.parameters(), lr=cfg.learning_rate) + self.buffer = ReplayBuffer(cfg.buffer_steps, transition) + self.error_floor = torch.cat( + ( + torch.full((3,), 0.003), + torch.full((3,), 0.005), + torch.full((3,), 0.02), + torch.full((3,), 0.05), + torch.full((29,), 0.003), + torch.full((29,), 0.05), + ) + ).to(self.device) + self.rollout = torch.compile(self._wm_rollout, disable=not cfg.compile) + self.error = torch.compile(self._wm_err, disable=not cfg.compile) + self.imagine = torch.compile(self._imag_step, disable=not cfg.compile) + self.reward_termination = torch.compile(self._reward_term, disable=not cfg.compile) + self.critic_input_compiled = torch.compile(self._critic_input, disable=not cfg.compile) + + def add_noise(self, observation): + return torch.cat( + [ + observation[..., region] if fn is None else fn(observation[..., region], cfg) + for region, fn, cfg in self.noise_terms + ], + -1, + ) + + @torch.no_grad() + def act(self, observations): + self.pi.obs_norm.update(observations["policy"][:, -1]) + self.pi.goal_norm.update(observations["goal"]) + action = self.pi(observations["policy"], observations["goal"]) + return (action + self.cfg.action_noise * torch.randn_like(action)).clamp(-20, 20) + + @torch.no_grad() + def observe(self, transition): + self.buffer.add(transition) + state, successor = (unflatten_state(transition["s"]), unflatten_state(transition["s1"])) + self.wm.in_norm.update(self.wm.inputs(state, transition["a"], normalize_xi(transition["xi"]))) + self.wm.out_norm.update(state_difference(state, successor)) + + def _wm_rollout(self, s, a, xi): + S = unflatten_state(s[:, 0]) + preds = [] + for j in range(self.cfg.world_model_window): + S = self.wm.step(S, a[:, j], normalize_xi(xi[:, j])) + preds.append(S) + return tuple(torch.stack([P[i] for P in preds], 1) for i in range(6)) + + def _wm_err(self, S_hat, s1, mask): + std = torch.maximum(self.wm.out_norm.std(), self.error_floor) + e = (state_difference(unflatten_state(s1), S_hat) / std) ** 2 + e = e.mean(-1) + return ((e * mask).sum() / s1.shape[0], e[:, -1]) + + def _imag_step(self, S, o_hist, a_prev, f, xi, xi_n, g): + a = self.pi(o_hist, g) + S1 = self.wm.step(S, a, xi_n) + f1 = f + 1 + o1 = self.add_noise(self.mdp.obs(S1, a, xi, f1)) + o_hist = torch.cat((o_hist[:, 1:], o1.unsqueeze(1)), 1) + return (S1, a, o_hist, f1, self.motion.goal(f1, self.goal_frames)) + + def _reward_term(self, S, f, a, a_prev): + return (self.reward_scale * self.mdp.reward(S, f, a, a_prev), self.mdp.soft_termination(S, f)) + + def _critic_input(self, S, f, xi_n, g): + return torch.cat((critic_features(self.mdp, S, f, xi_n), g), -1) + + def wm_update(self): + w = self.buffer.sample_windows(self.cfg.world_model_batch_size, self.cfg.world_model_window) + S_hat = self.rollout(w["s"], w["a"], w["xi"]) + loss, e_last = self.error(S_hat, w["s1"], w["mask"]) + self.opt_wm.zero_grad(set_to_none=True) + loss.backward() + nn.utils.clip_grad_norm_(self.wm.parameters(), 1.0, error_if_nonfinite=True) + self.opt_wm.step() + m = w["mask"][:, -1] + return (loss.detach() / self.cfg.world_model_window, (e_last.detach() * m).sum() / m.sum().clamp(min=1)) + + def sample_valid_starts(self, B: int): + starts, hists, n = ([], [], 0) + attempts = 0 + while n < B: + attempts += 1 + if attempts > 32: + raise RuntimeError("Replay has too few valid imagination starts for the configured horizon.") + s, h = self.buffer.sample_starts(2 * B, self.history) + f = s["f"] + ok = f + self.horizon <= self.motion.clip_end[self.motion.clip_id[f]] + idx = ok.nonzero(as_tuple=False).squeeze(-1)[: B - n] + starts.append({k: v[idx] for k, v in s.items()}) + hists.append({k: v[idx] for k, v in h.items()}) + n += len(idx) + return tuple( + {key: torch.cat([sample[key] for sample in samples]) for key in samples[0]} for samples in (starts, hists) + ) + + def policy_update(self): + s, h = self.sample_valid_starts(self.cfg.policy_batch_size) + S = unflatten_state(s["s"]) + xi, f, a_prev = (s["xi"], s["f"], s["a_prev"]) + xi_n = normalize_xi(xi) + with torch.no_grad(): + o_hist = self.mdp.obs(unflatten_state(h["s"]), h["a_prev"], h["xi"], h["f"]) + o_hist = self.add_noise(o_hist) + self.wm.requires_grad_(False) + self.critic.requires_grad_(False) + B = self.cfg.policy_batch_size + g0 = self.motion.goal(f, self.goal_frames) + g = g0 + S0 = S + Ss, As, Aps, Fs, Gs = ([], [], [], [], []) + for _ in range(self.horizon): + Aps.append(a_prev) + S, a_prev, o_hist, f, g = self.imagine(S, o_hist, a_prev, f, xi, xi_n, g) + Ss.append(S) + As.append(a_prev) + Fs.append(f) + Gs.append(g) + + def flat(xs): + return torch.stack(xs).flatten(0, 1) + + S_all = tuple(flat([St[i] for St in Ss]) for i in range(6)) + f_all = flat(Fs) + r, d = self.reward_termination(S_all, f_all, flat(As), flat(Aps)) + r, d = (r.view(self.horizon, B), d.view(self.horizon, B)) + alive = torch.cumprod(1 - d, 0) + c_t = torch.cat((torch.ones(1, B, device=self.device), alive[:-1]), 0) + disc = self.gamma ** torch.arange(self.horizon, device=self.device, dtype=torch.float32).unsqueeze(-1) + J = (disc * c_t * r).sum(0) + c = alive[-1] + r_sum = (c_t * r).detach().sum(0) + rews, dones = (r.detach(), d.detach()) + terminal_value = self.critic(self._critic_input(S, f, xi_n, g)) + J = J + self.gamma**self.horizon * c * terminal_value + loss = -J.mean() + self.opt_pi.zero_grad(set_to_none=True) + loss.backward() + gn = nn.utils.clip_grad_norm_(self.pi.parameters(), 1.0) + nn.utils.clip_grad_value_(self.pi.parameters(), 0.1) + if torch.isfinite(gn): + self.opt_pi.step() + self.wm.requires_grad_(True) + self.critic.requires_grad_(True) + with torch.no_grad(): + S_c = tuple( + torch.cat((S0[i].unsqueeze(0), torch.stack([St[i].detach() for St in Ss])), 0).flatten(0, 1) + for i in range(6) + ) + f_c = torch.cat((s["f"].unsqueeze(0), torch.stack(Fs)), 0).flatten() + g_c = torch.cat((g0.unsqueeze(0), torch.stack(Gs)), 0).flatten(0, 1) + xi_c = xi_n.unsqueeze(0).expand(self.horizon + 1, -1, -1).flatten(0, 1) + X = self.critic_input_compiled(S_c, f_c, xi_c, g_c).view(self.horizon + 1, B, -1) + self.critic.norm.update(X) + with torch.no_grad(): + V = self.critic(X) + R = V[-1] + targets = [] + for t in range(self.horizon - 1, -1, -1): + R = rews[t] + self.gamma * (1 - dones[t]) * ( + (1 - self.cfg.gae_lambda) * V[t + 1] + self.cfg.gae_lambda * R + ) + targets.append(R) + targets = torch.stack(targets[::-1]) + Xf, Yf = (X[:-1].reshape(-1, X.shape[-1]), targets.reshape(-1)) + n = Xf.shape[0] + mb = n // self.cfg.critic_minibatches + for _ in range(self.cfg.critic_epochs): + perm = torch.randperm(n, device=self.device) + for i in range(self.cfg.critic_minibatches): + idx = perm[i * mb : (i + 1) * mb] + prediction = self.critic(Xf[idx]) + v_loss = ((prediction - Yf[idx]) ** 2).mean() + self.opt_v.zero_grad(set_to_none=True) + v_loss.backward() + nn.utils.clip_grad_norm_(self.critic.parameters(), 1.0, error_if_nonfinite=True) + self.opt_v.step() + return ( + loss.detach(), + v_loss.detach(), + (r_sum / self.horizon / (MAX_REWARD_RATE * self.reward_scale)).mean(), + gn.detach(), + c.mean().detach(), + ) diff --git a/source/isaaclab_rl/isaaclab_rl/supermimic/networks.py b/source/isaaclab_rl/isaaclab_rl/supermimic/networks.py new file mode 100644 index 000000000000..1648ea39e2df --- /dev/null +++ b/source/isaaclab_rl/isaaclab_rl/supermimic/networks.py @@ -0,0 +1,165 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""World model, policy (motion encoder + shared decoder), and critic (paper Table A.4).""" + +from __future__ import annotations + +import torch +import torch.nn as nn +from physicsnemo_robotics.dynamics.normalization import RunningNormalizer + +from isaaclab_tasks.contrib.g1_tracking.kinematics import mat_to_6d, so3_exp, so3_log, yaw_mat +from isaaclab_tasks.contrib.g1_tracking.tracking import NUM_J, XI_DIM + +INV_DIM = 1 + 6 + 3 + 3 + 29 + 29 +CRITIC_EXTRA = 29 + 14 * 3 + 3 + 6 + XI_DIM + + +def invariant_features(S) -> torch.Tensor: + """Heading- and position-invariant state features: height, tilt (6D), body-frame velocities, joints.""" + p, R, v, w, q, qd = S + H = yaw_mat(R) + Rt = R.transpose(-1, -2) + return torch.cat( + ( + p[..., 2:3], + mat_to_6d(H.transpose(-1, -2) @ R), + (Rt @ v.unsqueeze(-1)).squeeze(-1), + (Rt @ w.unsqueeze(-1)).squeeze(-1), + q, + qd, + ), + -1, + ) + + +def normalize_xi(xi: torch.Tensor) -> torch.Tensor: + return torch.cat(((xi[..., :1] - 0.75) / 0.45, xi[..., 1:30] / 0.01, xi[..., 30:33] / 0.05), -1) + + +def mlp(inp: int, hidden: list[int], out: int, act=nn.ELU, layer_norm: bool = False) -> nn.Sequential: + layers, d = [], inp + for h in hidden: + layers.append(nn.Linear(d, h)) + if layer_norm: + layers.append(nn.LayerNorm(h)) + layers.append(act()) + d = h + layers.append(nn.Linear(d, out)) + return nn.Sequential(*layers) + + +WM_IN = INV_DIM + NUM_J + XI_DIM +WM_OUT = 3 + 3 + 3 + 3 + NUM_J + NUM_J + + +def run_network(net: nn.Module, x: torch.Tensor, dtype: torch.dtype | None) -> torch.Tensor: + if dtype is None: + return net(x) + with torch.autocast("cuda", dtype=dtype): + return net(x).float() + + +def state_difference(S, S1): + """Frame-local increments from S to S1 (the regression targets).""" + p, R, v, w, q, qd = S + p1, R1, v1, w1, q1, qd1 = S1 + H = yaw_mat(R) + Rt = R.transpose(-1, -2) + return torch.cat( + ( + (H.transpose(-1, -2) @ (p1 - p).unsqueeze(-1)).squeeze(-1), + so3_log(Rt @ R1), + (Rt @ (v1 - v).unsqueeze(-1)).squeeze(-1), + (Rt @ (w1 - w).unsqueeze(-1)).squeeze(-1), + q1 - q, + qd1 - qd, + ), + -1, + ) + + +class WorldModel(nn.Module): + """Deterministic dynamics s_{t+1} = f(s_t, a_t, xi) predicting frame-local state increments.""" + + def __init__(self, hidden=(2048, 2048, 2048), amp_dtype=None, network: nn.Module | None = None): + super().__init__() + self.amp_dtype = amp_dtype + self.net = network if network is not None else mlp(WM_IN, list(hidden), WM_OUT, act=nn.SiLU, layer_norm=True) + if network is None: + nn.init.zeros_(self.net[-1].weight) + nn.init.zeros_(self.net[-1].bias) + self.in_norm = RunningNormalizer(WM_IN) + self.out_norm = RunningNormalizer(WM_OUT, eps=1e-4) + + def inputs(self, S, a, xi_n): + return torch.cat((invariant_features(S), a, xi_n), -1) + + def step(self, S, a, xi_n): + p, R, v, w, q, qd = S + tokens = self.in_norm(self.inputs(S, a, xi_n)).unsqueeze(-2) + y = run_network(self.net, tokens, self.amp_dtype).squeeze(-2) + d = y * self.out_norm.std() + self.out_norm.mean + H = yaw_mat(R) + dp, dth, dv, dw, dq, dqd = d.split((3, 3, 3, 3, NUM_J, NUM_J), -1) + p1 = p + (H @ dp.unsqueeze(-1)).squeeze(-1) + R1 = R @ so3_exp(dth) + v1 = v + (R @ dv.unsqueeze(-1)).squeeze(-1) + w1 = w + (R @ dw.unsqueeze(-1)).squeeze(-1) + return (p1, R1, v1, w1, q + dq, qd + dqd) + + +class Policy(nn.Module): + """a = dec(enc(g), o_{t-9:t}); deterministic, latent z in R^32 bounded by tanh.""" + + def __init__(self, obs_dim: int, goal_dim: int, hist: int = 10, z_dim: int = 32, amp_dtype=None): + super().__init__() + self.hist = hist + self.amp_dtype = amp_dtype + self.enc = mlp(goal_dim, [2048, 1024, 512, 256], z_dim) + self.dec = mlp(z_dim + obs_dim * hist, [2048, 1024, 512, 256], NUM_J) + with torch.no_grad(): + self.dec[-1].weight.mul_(0.01) + self.dec[-1].bias.zero_() + self.obs_norm = RunningNormalizer(obs_dim) + self.goal_norm = RunningNormalizer(goal_dim) + + def forward(self, obs_hist: torch.Tensor, goal: torch.Tensor) -> torch.Tensor: + """obs_hist: (B, hist, obs_dim) noisy observations, newest last.""" + z = torch.tanh(run_network(self.enc, self.goal_norm(goal), self.amp_dtype)) + o = self.obs_norm(obs_hist).flatten(-2) + return run_network(self.dec, torch.cat((z, o), -1), self.amp_dtype) + + +class Critic(nn.Module): + def __init__(self, in_dim: int, amp_dtype=None): + super().__init__() + self.amp_dtype = amp_dtype + self.net = mlp(in_dim, [256, 128], 1) + self.norm = RunningNormalizer(in_dim) + + def forward(self, x): + return run_network(self.net, self.norm(x), self.amp_dtype).squeeze(-1) + + +def critic_features(mdp, S, f, xi) -> torch.Tensor: + """Critic input: invariant state, goal, sim context, and robot-vs-reference errors.""" + m = mdp.motion + p, R, v, w, q, qd = S + H = yaw_mat(R) + Ht = H.transpose(-1, -2) + rp, rR = m.root_p[f], m.root_R[f] + bp = mdp.kin.forward(p, R, v, w, q, qd)[0] + dR = yaw_mat(R @ rR.transpose(-1, -2)) + anchor = torch.cat((p[..., :2], rp[..., 2:3]), -1) + ref_bp = anchor.unsqueeze(-2) + ((dR.unsqueeze(-3) @ (m.body_p[f] - rp.unsqueeze(-2)).unsqueeze(-1)).squeeze(-1)) + e_body = (Ht.unsqueeze(-3) @ (ref_bp - bp).unsqueeze(-1)).squeeze(-1).flatten(-2) + e_root = (Ht @ (rp - p).unsqueeze(-1)).squeeze(-1) + rel = mat_to_6d(R.transpose(-1, -2) @ rR) + return torch.cat((invariant_features(S), m.q[f] - q, e_body, e_root, rel, xi), -1) diff --git a/source/isaaclab_rl/isaaclab_rl/supermimic/supermimic_cfg.py b/source/isaaclab_rl/isaaclab_rl/supermimic/supermimic_cfg.py new file mode 100644 index 000000000000..c5efffde2ee7 --- /dev/null +++ b/source/isaaclab_rl/isaaclab_rl/supermimic/supermimic_cfg.py @@ -0,0 +1,51 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Learning and checkpoint settings for the native SuperMimic workflow.""" + +from isaaclab.utils import configclass + + +@configclass +class SuperMimicCfg: + """SuperMimic settings; task timing, observations and reference goals belong to the environment.""" + + seed: int = 0 + experiment_name: str = "g1_tracking" + max_iterations: int = 14_000 + buffer_steps: int = 1000 + """Replay capacity in control steps per world.""" + update_to_data_ratio: float = 0.4 + """Learning iterations per physical control step.""" + horizon: int = 32 + """Differentiable imagined rollout length.""" + world_model_window: int = 8 + world_model_batches: int = 10 + world_model_batch_size: int = 2048 + policy_batch_size: int = 4096 + learning_rate: float = 2e-4 + gae_lambda: float = 0.95 + action_noise: float = 0.1 + critic_epochs: int = 5 + critic_minibatches: int = 4 + warmup_steps: int = 64 + world_model_only_iterations: int = 500 + learning_rate_decay_iterations: int = 250_000 + world_model: str = "mlp" + """Dynamics architecture: ``mlp`` or ``nerd-one-frame``. Both use PhysicsNeMo Robotics normalization.""" + amp: str = "bf16" + """Neural projection precision: ``bf16`` or ``none``. Geometry always uses float32.""" + compile: bool = True + """Compile the tensor rollout, objective and regression functions with Torch.""" + save_interval: int = 2000 + save_interval_seconds: float = 120.0 + """Maximum time between checkpoint requests, checked at iteration boundaries.""" + log_interval: int = 50 + max_hours: float | None = None + """Cumulative training-hour budget, including prior progress when resuming.""" + wandb_project: str | None = None + wandb_entity: str | None = None + wandb_run_id: str | None = None + wandb_name: str | None = None diff --git a/source/isaaclab_rl/test/test_entrypoints.py b/source/isaaclab_rl/test/test_entrypoints.py index 4f1f98ae4c13..92357c887ce5 100644 --- a/source/isaaclab_rl/test/test_entrypoints.py +++ b/source/isaaclab_rl/test/test_entrypoints.py @@ -46,6 +46,11 @@ ("import isaaclab_rl.rl_games.pbt", ["isaaclab_rl.rl_games.pbt.pbt", "rl_games", "torch"]), ("import isaaclab_rl.rsl_rl", ["isaaclab_rl.rsl_rl.vecenv_wrapper", "rsl_rl", "torch"]), ("import isaaclab_rl.utils", ["torch"]), + ( + "import isaaclab_tasks.contrib.g1_tracking", + ["isaaclab_rl.supermimic", "physicsnemo_robotics", "wandb", "supermimic"], + ), + ("import isaaclab_rl.supermimic.networks", ["physicsnemo_robotics", "wandb", "supermimic"]), # resolves through isaaclab_rl.entrypoints, so it also covers importing that package directly ( "from isaaclab_rl import run_play_cli", diff --git a/source/isaaclab_rl/test/test_supermimic.py b/source/isaaclab_rl/test/test_supermimic.py new file mode 100644 index 000000000000..8e487ababbdc --- /dev/null +++ b/source/isaaclab_rl/test/test_supermimic.py @@ -0,0 +1,231 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""Contracts for native SuperMimic ownership, replay and evaluation.""" + +import ast +from pathlib import Path + +import numpy as np +import pytest +import torch + +from isaaclab.paths import ISAACLAB_ROOT +from isaaclab.test.utils import DeviceScope, test_devices + +from isaaclab_rl.supermimic.learner import ReplayBuffer + + +@pytest.mark.parametrize("architecture", ("mlp", "nerd-one-frame")) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +def test_compiled_world_model_propagates_finite_action_gradients(architecture, device): + """Imagined control needs action gradients through frozen dynamics, including identity rotations.""" + from isaaclab_rl.supermimic.networks import WM_IN, WM_OUT, WorldModel, state_difference + + torch.manual_seed(0) + network = None + if architecture == "nerd-one-frame": + dependency = pytest.importorskip("physicsnemo_robotics.models.parallel_nerd") + network = dependency.ParallelNeRD(WM_IN, WM_OUT, temporal_attention_spans=(1, 1, 1, 1)) + torch.nn.init.normal_(network.head_output.weight, std=0.001) + world = WorldModel(hidden=(32, 32), network=network).to(device) + if network is None: + torch.nn.init.normal_(world.net[-1].weight, std=0.001) + world.requires_grad_(False) + state = ( + torch.zeros(4, 3, device=device), + torch.eye(3, device=device).expand(4, -1, -1), + torch.zeros(4, 3, device=device), + torch.zeros(4, 3, device=device), + torch.zeros(4, 29, device=device), + torch.zeros(4, 29, device=device), + ) + action = torch.zeros(4, 29, device=device, requires_grad=True) + context = torch.zeros(4, 33, device=device) + successor = torch.compile(world.step)(state, action, context) + loss = (state_difference(state, successor) - 0.1).square().mean() + gradient = torch.autograd.grad(loss, action)[0] + assert torch.isfinite(loss) and gradient.isfinite().all() and gradient.abs().sum() > 0 + assert all(parameter.grad is None for parameter in world.parameters()) + + +def test_native_ownership_has_no_external_training_package(): + """Reject external forwarders, duplicate task stepping and task-to-learner dependency inversions.""" + task = Path(ISAACLAB_ROOT) / "source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking" + learning = Path(ISAACLAB_ROOT) / "source/isaaclab_rl/isaaclab_rl/supermimic" + backends = learning.parent / "entrypoints/backends" + roots = (backends / "train_supermimic.py", backends / "play_supermimic.py") + for path in (*task.rglob("*.py"), *learning.rglob("*.py"), *roots): + tree = ast.parse(path.read_text()) + imports = [] + for node in ast.walk(tree): + if isinstance(node, ast.Import): + imports.extend(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.level == 0: + imports.append(node.module or "") + assert not any(name == "supermimic" or name.startswith("supermimic.") for name in imports), path + symbols = {node.name for node in ast.walk(tree) if isinstance(node, (ast.ClassDef, ast.FunctionDef))} + assert symbols.isdisjoint({"TrackingSim", "TransitionCapture", "RunningNorm", "RunningNormalizer"}), path + if path.is_relative_to(task): + assert not any(name.startswith(("isaaclab_rl", "physicsnemo", "wandb")) for name in imports), path + assert "step" not in symbols and path.name not in ("train.py", "learner.py"), path + elif path.is_relative_to(learning): + assert path.name not in ("train.py", "evaluate.py", "sim.py", "robot.py"), path + assert not any(name.startswith(("argparse", "gymnasium", "wandb", "isaaclab.envs")) for name in imports), ( + path + ) + + +def test_native_task_registration_and_optional_capture(): + from isaaclab_tasks.utils import resolve_task_config + + env_cfg, agent_cfg = resolve_task_config( + "IsaacContrib-Motion-Tracking-G1", "supermimic_cfg_entry_point", overrides=[] + ) + assert env_cfg.recorders.transition is None + assert type(env_cfg).__module__.startswith("isaaclab_tasks.") + assert type(agent_cfg).__module__.startswith("isaaclab_rl.") + + +def test_replay_ring_masks_boundaries_and_owns_samples(): + transition = { + "s": torch.zeros(1, 71), + "s1": torch.zeros(1, 71), + "a": torch.zeros(1, 29), + "a_prev": torch.zeros(1, 29), + "xi": torch.zeros(1, 33), + "r": torch.zeros(1), + "f": torch.zeros(1, dtype=torch.long), + "ep_step": torch.zeros(1, dtype=torch.long), + "done": torch.zeros(1, dtype=torch.bool), + } + ring = ReplayBuffer(4, transition) + assert "r" not in ring.data # This learner evaluates analytical rewards; capture still supplies them to others. + for i in range(6): + transition["s"].fill_(i) + transition["s1"].fill_(i + 1) + transition["f"].fill_(i) + transition["ep_step"].fill_(i) + transition["done"].fill_(i == 3) + ring.add(transition) + transition["s"].fill_(99) + window = ring.sample_windows(4, 4) + torch.testing.assert_close(window["s"][..., 0], torch.tensor([[2.0, 3.0, 4.0, 5.0]]).expand(4, -1)) + torch.testing.assert_close(window["mask"], torch.tensor([[1.0, 1.0, 0.0, 0.0]]).expand(4, -1)) + # A state intervention or manual reset is a boundary even when no sampled done flag was set. + ring.data["done"].zero_() + ring.data["s"][ring._row(1), 0, 0] = -1 + torch.testing.assert_close(ring.sample_windows(4, 4)["mask"], torch.tensor([[1.0, 0.0, 0.0, 0.0]]).expand(4, -1)) + ring.data["ep_step"].zero_() + start, history = ring.sample_starts(32, 10) + torch.testing.assert_close(history["s"], start["s"][:, None].expand(-1, 10, -1)) + + +def test_evaluation_batches_cover_last_partial_batch_and_failed_tails(): + from types import SimpleNamespace + + from isaaclab_rl.entrypoints.backends.play_supermimic import evaluate + + lengths = torch.tensor([3, 4, 2, 5, 3, 4, 2]) + motion = SimpleNamespace(names=[f"clip{i}" for i in range(7)], clip_len=lengths) + command = SimpleNamespace(motion=motion, cfg=SimpleNamespace(evaluation_clip_offset=0)) + capture = SimpleNamespace(latest={}) + + class Env: + num_envs, device, step_dt = 3, "cpu", 0.02 + command_manager = SimpleNamespace(get_term=lambda name: command) + recorder_manager = SimpleNamespace(get_term=lambda name: capture) + reward_manager = SimpleNamespace(get_term_cfg=lambda name: SimpleNamespace(weight=1.0)) + + @property + def unwrapped(self): + return self + + def reset(self): + self.ids = (torch.arange(3) + command.cfg.evaluation_clip_offset) % 7 + self.elapsed = 0 + return {"policy": None, "goal": None}, {} + + def step(self, action): + self.elapsed += 1 + term = (self.ids == 3) & (self.elapsed == 1) + capture.latest = { + "r": (self.ids + 1) * 0.01, + "done": term | (self.elapsed >= lengths[self.ids] - 1), + "term": term, + } + return {"policy": None, "goal": None}, None, None, None, None + + result = evaluate(Env(), lambda *args: None) + assert result["failed"] == ["clip3"] + assert list(result["per_clip_steps"].values()) == [2, 3, 1, 1, 2, 3, 1] + np.testing.assert_allclose(list(result["per_clip"].values()), [0.1, 0.2, 0.3, 0.1, 0.5, 0.6, 0.7], rtol=1e-6) + + +def test_training_root_rejects_different_analytical_task_semantics(tmp_path, monkeypatch): + from isaaclab.envs import mdp + + from isaaclab_rl.entrypoints.backends import train_supermimic as train + from isaaclab_rl.supermimic.supermimic_cfg import SuperMimicCfg + + from isaaclab_tasks.contrib.g1_tracking.g1_tracking_env_cfg import G1TrackingEnvCfg + + class ReachedLauncher(Exception): + pass + + def launch(*args): + raise ReachedLauncher + + monkeypatch.setattr(train, "launch_simulation", launch) + invalid = ( + (("policy", "joint_vel", "func"), mdp.joint_pos), # Same width, different physical quantity. + (("policy", "joint_vel", "params"), {"asset_cfg": "different_asset"}), + (("policy", "joint_vel", "delay_max_lag"), 1), + (("policy", "flatten_history_dim"), True), + (("goal", "reference", "scale"), 2.0), + ) + argv = ["--task", "IsaacContrib-Motion-Tracking-G1", "--log_dir", str(tmp_path)] + for path, value in invalid: + cfg = G1TrackingEnvCfg() + target = cfg.observations + for name in path[:-1]: + target = getattr(target, name) + setattr(target, path[-1], value) + monkeypatch.setattr(train, "resolve_task_config", lambda *args, **kwargs: (cfg, SuperMimicCfg())) + with pytest.raises(ValueError, match="declared functions and layout"): + train.run(argv) + cfg = G1TrackingEnvCfg() + term = cfg.observations.policy.__dict__.pop("joint_vel") + cfg.observations.policy.joint_vel = term + monkeypatch.setattr(train, "resolve_task_config", lambda *args, **kwargs: (cfg, SuperMimicCfg())) + with pytest.raises(ValueError, match="declared functions and layout"): + train.run(argv) + for name, attribute, value in ( + ("failed", "func", mdp.time_out), + ("failed", "time_out", True), + ("clip_end", None, None), + ): + cfg = G1TrackingEnvCfg() + if attribute is None: + setattr(cfg.terminations, name, value) + else: + setattr(getattr(cfg.terminations, name), attribute, value) + monkeypatch.setattr(train, "resolve_task_config", lambda *args, **kwargs: (cfg, SuperMimicCfg())) + with pytest.raises(ValueError, match="Analytical termination"): + train.run(argv) + cfg = G1TrackingEnvCfg() + cfg.observations.policy.enable_corruption = False + cfg.observations.policy.joint_vel.noise.n_max = 0.25 + cfg.commands.motion.objective.term_height = 0.3 + cfg.episode_length_s = 4.0 + monkeypatch.setattr(train, "resolve_task_config", lambda *args, **kwargs: (cfg, SuperMimicCfg())) + with pytest.raises(ReachedLauncher): + train.run(argv) # Supported task parameters reach construction without a simulator. + cfg.observations.policy.history_length = 0 + with pytest.raises(ValueError, match="must be positive"): + train.run(argv) diff --git a/source/isaaclab_tasks/changelog.d/g1-tracking.added.rst b/source/isaaclab_tasks/changelog.d/g1-tracking.added.rst new file mode 100644 index 000000000000..3851fd1d62ec --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/g1-tracking.added.rst @@ -0,0 +1,2 @@ +* Added the contributed G1 motion-tracking task with native manager terms, optional transition capture, + differentiable tracking objectives and stationary goal padding beyond recorded clip endpoints. diff --git a/source/isaaclab_tasks/changelog.d/g1-tracking.minor b/source/isaaclab_tasks/changelog.d/g1-tracking.minor new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_tasks/changelog.d/native-nerd.added.rst b/source/isaaclab_tasks/changelog.d/native-nerd.added.rst new file mode 100644 index 000000000000..b37935b2f6f3 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/native-nerd.added.rst @@ -0,0 +1,7 @@ +* Added native Cartpole and ANYmal C NeRD exploration, numerical-teacher and learned-physics + task configurations. Policies use the standard reinforcement-learning training and playback scripts. + ANYmal keeps ANYdrive outside NeRD and models each 5 ms physics step from applied effort. + Cartpole's actuator owns damping and effort limits; neither task uses solver-side motor drives. +* Allowed velocity task playback with an explicit ground asset instead of a terrain importer. +* Added rough-terrain ANYmal C numerical, exploration and learned tasks using live Newton contacts, + terrain-aware reset sampling and a contact-conditioned NeRD fitting recipe. diff --git a/source/isaaclab_tasks/changelog.d/native-nerd.minor b/source/isaaclab_tasks/changelog.d/native-nerd.minor new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_tasks/changelog.d/solver-config-replacement.fixed.rst b/source/isaaclab_tasks/changelog.d/solver-config-replacement.fixed.rst new file mode 100644 index 000000000000..26cade37fe0e --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/solver-config-replacement.fixed.rst @@ -0,0 +1 @@ +* Removed the Digit preset workaround for copying the derived Newton manager class. diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/__init__.py new file mode 100644 index 000000000000..513fa5cf146d --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/__init__.py @@ -0,0 +1,19 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""G1 reference-motion tracking with native environment managers.""" + +import gymnasium as gym + +gym.register( + id="IsaacContrib-Motion-Tracking-G1", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.g1_tracking_env_cfg:G1TrackingEnvCfg", + "default_agent": "supermimic", + "supermimic_cfg_entry_point": "isaaclab_rl.supermimic.supermimic_cfg:SuperMimicCfg", + }, +) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/g1_tracking_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/g1_tracking_env_cfg.py new file mode 100644 index 000000000000..905194202ab9 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/g1_tracking_env_cfg.py @@ -0,0 +1,250 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""The G1 tracking task's native scene, manager terms and Newton physics preset.""" + +from pathlib import Path + +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +from isaaclab_newton.sim.schemas import NewtonArticulationCfg +from isaaclab_physx.sim.schemas import PhysxRigidBodyCfg + +import isaaclab.sim as sim_utils +from isaaclab.actuators import ImplicitActuatorCfg +from isaaclab.assets import ArticulationCfg, AssetBaseCfg +from isaaclab.envs import ManagerBasedRLEnvCfg, mdp +from isaaclab.managers import ( + CommandTermCfg, + DatasetExportMode, + EventTermCfg, + ObservationGroupCfg, + ObservationTermCfg, + RecorderManagerBaseCfg, + RecorderTermCfg, + RewardTermCfg, + TerminationTermCfg, +) +from isaaclab.paths import ISAACLAB_ROOT +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.utils import configclass, replace +from isaaclab.utils.noise import UniformNoiseCfg + +from isaaclab_tasks.utils import PresetCfg + +from .motion import CONTROL_DT +from .tracking import RewardCfg + +G1_DATA_DIR = Path(ISAACLAB_ROOT) / "data/g1_tracking" + +ARMATURE_5020 = 0.003609725 +ARMATURE_7520_14 = 0.010177520 +ARMATURE_7520_22 = 0.025101925 +ARMATURE_4010 = 0.00425 + +NATURAL_FREQ = 10 * 2.0 * 3.1415926535 # 10 Hz +DAMPING_RATIO = 2.0 + + +# Per-joint (effort limit, velocity limit, armature) by motor type. +_M5020 = (25.0, 37.0, ARMATURE_5020) +_M4010 = (5.0, 22.0, ARMATURE_4010) +JOINT_MOTORS = { + "hip_pitch": (88.0, 32.0, ARMATURE_7520_14), + "hip_roll": (139.0, 20.0, ARMATURE_7520_22), + "hip_yaw": (88.0, 32.0, ARMATURE_7520_14), + "knee": (139.0, 20.0, ARMATURE_7520_22), + "ankle_pitch": (50.0, 37.0, 2.0 * ARMATURE_5020), + "ankle_roll": (50.0, 37.0, 2.0 * ARMATURE_5020), + "waist_yaw": (88.0, 32.0, ARMATURE_7520_14), + "waist_roll": (50.0, 37.0, 2.0 * ARMATURE_5020), + "waist_pitch": (50.0, 37.0, 2.0 * ARMATURE_5020), + "shoulder_pitch": _M5020, + "shoulder_roll": _M5020, + "shoulder_yaw": _M5020, + "elbow": _M5020, + "wrist_roll": _M5020, + "wrist_pitch": _M4010, + "wrist_yaw": _M4010, +} + + +DEFAULT_JOINT_POS = { + ".*_hip_pitch_joint": -0.312, + ".*_knee_joint": 0.669, + ".*_ankle_pitch_joint": -0.363, + ".*_elbow_joint": 0.6, + "left_shoulder_roll_joint": 0.2, + "left_shoulder_pitch_joint": 0.2, + "right_shoulder_roll_joint": -0.2, + "right_shoulder_pitch_joint": 0.2, +} + + +G1_CFG = ArticulationCfg( + # Prepare the matching 29-DoF MJCF and convert it to USD before constructing this task. + spawn=sim_utils.UsdFileCfg( + usd_path=str(G1_DATA_DIR / "g1_usd/g1_29dof/g1_29dof.usda"), + variants={"Physics": "mujoco"}, + activate_contact_sensors=False, + rigid_props=PhysxRigidBodyCfg(disable_gravity=False, linear_damping=0.0, angular_damping=0.0), + articulation_props=NewtonArticulationCfg(self_collision_enabled=False), + ), + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.0, 0.0, 0.76), joint_pos=DEFAULT_JOINT_POS, joint_vel={".*": 0.0} + ), + soft_joint_pos_limit_factor=0.9, + actuators={ + "joints": ImplicitActuatorCfg( + joint_names_expr=[".*"], + joint_effort_limit={f".*{name}_joint": motor[0] for name, motor in JOINT_MOTORS.items()}, + joint_velocity_limit={f".*{name}_joint": motor[1] for name, motor in JOINT_MOTORS.items()}, + stiffness={f".*{name}_joint": motor[2] * NATURAL_FREQ**2 for name, motor in JOINT_MOTORS.items()}, + damping={ + f".*{name}_joint": 2 * DAMPING_RATIO * motor[2] * NATURAL_FREQ for name, motor in JOINT_MOTORS.items() + }, + armature={f".*{name}_joint": motor[2] for name, motor in JOINT_MOTORS.items()}, + ) + }, +) + + +@configclass +class TrackingPhysicsCfg(PresetCfg): + newton_mjwarp = NewtonCfg( + solver_cfg=MJWarpSolverCfg( + njmax=500, nconmax=128, cone="pyramidal", integrator="implicitfast", use_mujoco_contacts=True + ), + num_substeps=4, + use_cuda_graph=True, + ) + default = newton_mjwarp + + +@configclass +class TrackingSceneCfg(InteractiveSceneCfg): + ground = AssetBaseCfg( + prim_path="/World/ground", + spawn=sim_utils.GroundPlaneCfg( + physics_material=sim_utils.RigidBodyMaterialBaseCfg(static_friction=0.3, dynamic_friction=0.3) + ), + ) + robot = replace(G1_CFG, prim_path="{ENV_REGEX_NS}/Robot") + + +@configclass +class MotionCommandCfg(CommandTermCfg): + class_type = "isaaclab_tasks.contrib.g1_tracking.mdp:MotionCommand" + resampling_time_range = (1.0e9, 1.0e9) + motion_dir: str = str(G1_DATA_DIR / "lafan1_g1") + """Directory containing metric/xyzw/radian 36-channel motion CSV or NumPy files.""" + kinematics_path: str = str(G1_DATA_DIR / "g1_description/g1_29dof.xml") + """MJCF matching the scene asset, used for differentiable link kinematics.""" + source_fps: float = 30.0 + goal_frames: int = 10 + gamma: float = 0.99 + evaluation: bool = False + evaluation_clip_offset: int = 0 + reset_noise: bool = True + objective: RewardCfg = RewardCfg() + + +@configclass +class CommandsCfg: + motion = MotionCommandCfg() + + +@configclass +class ActionsCfg: + joint_pos = mdp.JointPositionActionCfg( + class_type="isaaclab_tasks.contrib.g1_tracking.mdp:ContextJointPositionAction", + asset_name="robot", + joint_names=[".*"], + scale={ + name: 0.25 * G1_CFG.actuators["joints"].joint_effort_limit[name] / stiffness + for name, stiffness in G1_CFG.actuators["joints"].stiffness.items() + }, + ) + + +@configclass +class ObservationsCfg: + @configclass + class PolicyCfg(ObservationGroupCfg): + joint_pos = ObservationTermCfg( + func="isaaclab_tasks.contrib.g1_tracking.mdp:joint_pos_rel", noise=UniformNoiseCfg(n_min=-0.01, n_max=0.01) + ) + joint_vel = ObservationTermCfg(func=mdp.joint_vel, noise=UniformNoiseCfg(n_min=-0.5, n_max=0.5)) + angular_velocity = ObservationTermCfg(func=mdp.base_ang_vel, noise=UniformNoiseCfg(n_min=-0.2, n_max=0.2)) + gravity = ObservationTermCfg(func=mdp.projected_gravity, noise=UniformNoiseCfg(n_min=-0.05, n_max=0.05)) + previous_action = ObservationTermCfg(func=mdp.last_action) + orientation = ObservationTermCfg( + func="isaaclab_tasks.contrib.g1_tracking.mdp:relative_orientation", + noise=UniformNoiseCfg(n_min=-0.05, n_max=0.05), + ) + history_length = 10 + flatten_history_dim = False + concatenate_terms = True + enable_corruption = True + + @configclass + class GoalCfg(ObservationGroupCfg): + reference = ObservationTermCfg(func=mdp.generated_commands, params={"command_name": "motion"}) + + policy = PolicyCfg() + goal = GoalCfg() + + +@configclass +class RewardsCfg: + tracking = RewardTermCfg(func="isaaclab_tasks.contrib.g1_tracking.mdp:tracking_reward", weight=1.0) + + +@configclass +class TerminationsCfg: + failed = TerminationTermCfg(func="isaaclab_tasks.contrib.g1_tracking.mdp:tracking_failure") + clip_end = TerminationTermCfg(func="isaaclab_tasks.contrib.g1_tracking.mdp:clip_end", time_out=True) + time_out = TerminationTermCfg(func=mdp.time_out, time_out=True) + + +@configclass +class EventsCfg: + context = EventTermCfg(func="isaaclab_tasks.contrib.g1_tracking.mdp:randomize_context", mode="startup") + + +@configclass +class TransitionRecorderCfg(RecorderTermCfg): + class_type = "isaaclab_tasks.contrib.g1_tracking.mdp:TransitionRecorder" + record_to_episodes: bool = False + + +@configclass +class RecordersCfg(RecorderManagerBaseCfg): + transition: TransitionRecorderCfg | None = None + dataset_export_mode = DatasetExportMode.EXPORT_NONE + + +@configclass +class G1TrackingEnvCfg(ManagerBasedRLEnvCfg): + scene = TrackingSceneCfg(num_envs=4096, env_spacing=0.0) + commands = CommandsCfg() + actions = ActionsCfg() + observations = ObservationsCfg() + rewards = RewardsCfg() + terminations = TerminationsCfg() + events = EventsCfg() + recorders = RecordersCfg() + sim = sim_utils.SimulationCfg(dt=CONTROL_DT, render_interval=1, physics=TrackingPhysicsCfg()) + decimation = 1 + episode_length_s = 10.0 + + def play_mode(self): + self.commands.motion.evaluation = True + self.observations.policy.enable_corruption = False + self.episode_length_s = 1.0e9 + self.scene.ground.spawn.physics_material.static_friction = 1.0 + self.scene.ground.spawn.physics_material.dynamic_friction = 1.0 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/kinematics.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/kinematics.py new file mode 100644 index 000000000000..9988b66f3fdc --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/kinematics.py @@ -0,0 +1,164 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""Differentiable G1 geometry shared by reference motion, real states and imagined dynamics. + +Native math owns quaternion conversion (xyzw). The smooth SO(3) exponential/logarithm +and regularized heading projection retain the derivatives required by the learned rollout. +""" + +from __future__ import annotations + +import mujoco +import numpy as np +import torch + +from isaaclab.utils.math import axis_angle_from_quat, matrix_from_quat, quat_from_angle_axis, quat_from_matrix + +# Joint order of the LAFAN1 retarget CSVs (Unitree SDK order). +G1_JOINT_NAMES = [ + "left_hip_pitch_joint", "left_hip_roll_joint", "left_hip_yaw_joint", "left_knee_joint", + "left_ankle_pitch_joint", "left_ankle_roll_joint", + "right_hip_pitch_joint", "right_hip_roll_joint", "right_hip_yaw_joint", "right_knee_joint", + "right_ankle_pitch_joint", "right_ankle_roll_joint", + "waist_yaw_joint", "waist_roll_joint", "waist_pitch_joint", + "left_shoulder_pitch_joint", "left_shoulder_roll_joint", "left_shoulder_yaw_joint", "left_elbow_joint", + "left_wrist_roll_joint", "left_wrist_pitch_joint", "left_wrist_yaw_joint", + "right_shoulder_pitch_joint", "right_shoulder_roll_joint", "right_shoulder_yaw_joint", "right_elbow_joint", + "right_wrist_roll_joint", "right_wrist_pitch_joint", "right_wrist_yaw_joint", +] # fmt: skip + +# BeyondMimic's tracked bodies; the anchor/root is the pelvis. +TRACKED_BODIES = [ + "pelvis", "left_hip_roll_link", "left_knee_link", "left_ankle_roll_link", + "right_hip_roll_link", "right_knee_link", "right_ankle_roll_link", "torso_link", + "left_shoulder_roll_link", "left_elbow_link", "left_wrist_yaw_link", + "right_shoulder_roll_link", "right_elbow_link", "right_wrist_yaw_link", +] # fmt: skip + + +def so3_exp(v: torch.Tensor) -> torch.Tensor: + """Rotation vector (..., 3) -> matrix, smooth at zero.""" + theta2 = (v * v).sum(-1, keepdim=True) + theta = torch.sqrt(theta2 + 1e-12) + small = theta2 < 1e-8 + a = torch.where(small, 1 - theta2 / 6, torch.sin(theta) / theta) + b = torch.where(small, 0.5 - theta2 / 24, (1 - torch.cos(theta)) / (theta2 + 1e-12)) + x, y, z = v.unbind(-1) + zero = torch.zeros_like(x) + K = torch.stack((zero, -z, y, z, zero, -x, -y, x, zero), -1).reshape(v.shape[:-1] + (3, 3)) + eye = torch.eye(3, device=v.device, dtype=v.dtype).expand(K.shape) + return eye + a[..., None] * K + b[..., None] * (K @ K) + + +def so3_log(R: torch.Tensor) -> torch.Tensor: + """Matrix to rotation vector with a finite zero-angle derivative for model fitting.""" + return axis_angle_from_quat(quat_from_matrix(R)) + + +def rot_angle_sq(Ra: torch.Tensor, Rb: torch.Tensor) -> torch.Tensor: + """Squared geodesic angle between rotations, differentiable (via acos of trace, clamped).""" + tr = (Ra * Rb).sum((-1, -2)) # trace(Ra^T Rb) + cos = torch.clamp((tr - 1) / 2, -1 + 1e-6, 1 - 1e-6) + return torch.acos(cos) ** 2 + + +def yaw_mat(R: torch.Tensor) -> torch.Tensor: + """Heading (yaw-only) rotation of R, from the projected x axis.""" + fx, fy = R[..., 0, 0], R[..., 1, 0] + n = torch.sqrt(fx * fx + fy * fy + 1e-9) + c, s = fx / n, fy / n + o, one = torch.zeros_like(c), torch.ones_like(c) + return torch.stack((c, -s, o, s, c, o, o, o, one), -1).reshape(R.shape[:-2] + (3, 3)) + + +def mat_to_6d(R: torch.Tensor) -> torch.Tensor: + """First two columns, flattened (6D continuous rotation representation).""" + return R[..., :, :2].reshape(R.shape[:-2] + (6,)) + + +class G1Kinematics: + def __init__(self, mjcf_path: str, device: str, tracked: list[str] = TRACKED_BODIES): + m = mujoco.MjModel.from_xml_path(mjcf_path) + self.device = device + jname_to_idx = {n: i for i, n in enumerate(G1_JOINT_NAMES)} + # moving bodies in tree order (mujoco orders parents before children) + bodies, parent, joint_idx, axes, offs_p, offs_R = [], [], [], [], [], [] + mj_to_local = {} + for b in range(1, m.nbody): + name = m.body(b).name + if m.body_jntnum[b] == 0: + continue # fixed links are merged visually only; none are tracked + jid = m.body_jntadr[b] + mj_to_local[b] = len(bodies) + bodies.append(name) + if m.jnt_type[jid] == mujoco.mjtJoint.mjJNT_FREE: + parent.append(-1) + joint_idx.append(-1) + else: + pb = m.body_parentid[b] + if pb not in mj_to_local: + raise RuntimeError(f"fixed parent for {name}") + parent.append(mj_to_local[pb]) + joint_idx.append(jname_to_idx[m.joint(jid).name]) + axes.append(m.jnt_axis[jid].copy()) + offs_p.append(m.body_pos[b].copy()) + offs_R.append(m.body_quat[b].copy()) + self.body_names = bodies + self.parent = parent + self.joint_idx = joint_idx + self.axes = torch.tensor(np.array(axes), dtype=torch.float32, device=device) + self.off_p = torch.tensor(np.array(offs_p), dtype=torch.float32, device=device) + self.off_R = matrix_from_quat( + torch.tensor(np.array(offs_R)[:, [1, 2, 3, 0]], dtype=torch.float32, device=device) + ) + self.tracked_idx = [bodies.index(n) for n in tracked] + self.num_bodies = len(bodies) + # joint limits (in G1_JOINT_NAMES order) + lo, hi = np.zeros(29), np.zeros(29) + for j in range(m.njnt): + n = m.joint(j).name + if n in jname_to_idx: + lo[jname_to_idx[n]], hi[jname_to_idx[n]] = m.jnt_range[j] + self.q_lo = torch.tensor(lo, dtype=torch.float32, device=device) + self.q_hi = torch.tensor(hi, dtype=torch.float32, device=device) + + def forward(self, p, R, v, w, q, qd, all_bodies: bool = False): + """Tracked-body world positions, rotations, linear (link-origin) and angular velocities. + + Args: + p, R, v, w: root (pelvis) position (B,3), rotation (B,3,3), world lin/ang velocity (B,3). + q, qd: joint positions/velocities (B,29) in G1_JOINT_NAMES order. + """ + P, Rs, V, W = ( + [None] * self.num_bodies, + [None] * self.num_bodies, + [None] * self.num_bodies, + [None] * self.num_bodies, + ) + for i in range(self.num_bodies): + par = self.parent[i] + if par < 0: + P[i], Rs[i], V[i], W[i] = p, R, v, w + continue + j = self.joint_idx[i] + Rp = Rs[par] + pi = P[par] + (Rp @ self.off_p[i].unsqueeze(-1)).squeeze(-1) + R_off = Rp @ self.off_R[i] + Ri = R_off @ matrix_from_quat(quat_from_angle_axis(q[..., j], self.axes[i])) + axis_w = (R_off @ self.axes[i].unsqueeze(-1)).squeeze(-1) + wi = W[par] + axis_w * qd[..., j : j + 1] + vi = V[par] + torch.cross(W[par], pi - P[par], dim=-1) + P[i], Rs[i], V[i], W[i] = pi, Ri, vi, wi + idx = range(self.num_bodies) if all_bodies else self.tracked_idx + return ( + torch.stack([P[i] for i in idx], -2), + torch.stack([Rs[i] for i in idx], -3), + torch.stack([V[i] for i in idx], -2), + torch.stack([W[i] for i in idx], -2), + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/mdp.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/mdp.py new file mode 100644 index 000000000000..65bd9e98d0f4 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/mdp.py @@ -0,0 +1,276 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""Native manager terms for reference tracking and physical transition capture.""" + +from __future__ import annotations + +import torch +import warp as wp +from newton import ModelFlags + +from isaaclab.envs.mdp.actions import JointPositionAction +from isaaclab.managers import CommandTerm, RecorderTerm +from isaaclab.utils.math import matrix_from_quat, quat_from_matrix + +from .kinematics import G1_JOINT_NAMES, G1Kinematics, mat_to_6d, so3_exp +from .motion import BIN_FRAMES, MotionLib +from .tracking import MAX_REWARD_RATE, STATE_DIM, TrackingObjective + + +class MotionCommand(CommandTerm): + """Own reference sampling, physical context and the task's analytical objectives.""" + + def __init__(self, cfg, env): + super().__init__(cfg, env) + self.robot = env.scene["robot"] + if list(self.robot.joint_names) != G1_JOINT_NAMES: + raise ValueError(f"The tracking asset must preserve the motion joint order: {self.robot.joint_names}") + self.kin = G1Kinematics(cfg.kinematics_path, self.device) + self.motion = MotionLib(cfg.motion_dir, self.kin, self.device, cfg.source_fps) + self.mdp = TrackingObjective( + self.kin, self.motion, self.robot.data.default_joint_pos.torch[0].clone(), cfg.objective + ) + self.physics = env.sim.backend_manager + self.robot_view = self.robot.root_view + self.torso_id = self.robot.find_bodies("torso_link")[0][0] + self.env_ids = torch.arange(self.num_envs, device=self.device) + self.xi = torch.zeros(self.num_envs, 33, device=self.device) + self.start_frame = torch.zeros(self.num_envs, dtype=torch.long, device=self.device) + self.start_bin = torch.zeros_like(self.start_frame) + self._accounted_step = torch.zeros_like(self.start_frame) + self.episode_return = torch.zeros(self.num_envs, device=self.device) + self.discounted_return = torch.zeros_like(self.episode_return) + self.bin_return = torch.zeros(self.motion.num_bins, device=self.device) + self.bin_weight = torch.full_like(self.bin_return, 1.0 / self.motion.num_bins) + self._command = torch.zeros(self.num_envs, 73 * cfg.goal_frames, device=self.device) + self.metrics = { + "return": torch.zeros_like(self.episode_return), + "length": torch.zeros_like(self.episode_return), + } + + @property + def command(self): + return self._command + + @property + def frame(self): + return self.start_frame + self._env.episode_length_buf + + def state(self): + data = self.robot.data + pose, velocity = data.root_link_pose_w.torch, data.root_link_vel_w.torch + return ( + pose[:, :3] - self._env.scene.env_origins, + matrix_from_quat(pose[:, 3:7]), + velocity[:, :3], + velocity[:, 3:], + data.joint_pos.torch, + data.joint_vel.torch, + ) + + def _account_reward(self, env_ids): + steps = self._env.episode_length_buf[env_ids] + pending = steps > self._accounted_step[env_ids] + if not hasattr(self._env, "reward_buf"): + return pending + reward = torch.where(pending, self._env.reward_buf[env_ids], 0.0) + self.episode_return[env_ids] += reward + self.discounted_return[env_ids] += self.cfg.gamma ** (steps.float() - 1) * reward + self._accounted_step[env_ids] = steps + return pending + + def reset(self, env_ids=None): + ids = self.env_ids if env_ids is None else self.env_ids[env_ids] + pending = self._account_reward(ids) + # Public reset after command.compute discards a partial episode without updating its bin twice. + if hasattr(self._env, "reset_terminated") and not self.cfg.evaluation: + boundary = pending & (self._env.reset_terminated[ids] | self._env.reset_time_outs[ids]) + ended = ids[boundary] + if len(ended): + steps = self._env.episode_length_buf[ended].float() + bootstrap = torch.where( + self._env.reset_terminated[ended], + 0.0, + self.cfg.gamma**steps * (self.episode_return[ended] / steps) / (1 - self.cfg.gamma), + ) + returns = self.discounted_return[ended] + bootstrap + sums, counts = torch.zeros_like(self.bin_return), torch.zeros_like(self.bin_return) + sums.index_add_(0, self.start_bin[ended], returns) + counts.index_add_(0, self.start_bin[ended], torch.ones_like(returns)) + self.bin_return.copy_( + torch.where(counts > 0, 0.8 * self.bin_return + 0.2 * sums / counts.clamp_min(1), self.bin_return) + ) + scale = MAX_REWARD_RATE * self._env.step_dt * self._env.reward_manager.get_term_cfg("tracking").weight + self.metrics["return"][ids] = self.episode_return[ids] / ( + self._env.episode_length_buf[ids].clamp_min(1) * scale + ) + self.metrics["length"][ids] = self._env.episode_length_buf[ids].float() + return super().reset(ids) + + def _update_metrics(self): + self._account_reward(self.env_ids) + + def _update_command(self): + self._command.copy_(self.motion.goal(self.frame, self.cfg.goal_frames)) + + def _resample_command(self, env_ids): + ids = self.env_ids[env_ids] + n, motion = len(ids), self.motion + if self.cfg.evaluation: + frames = motion.clip_start[(ids + self.cfg.evaluation_clip_offset) % len(motion.names)] + bins = motion.frame_bin[frames] + else: + maximum_return = ( + MAX_REWARD_RATE + * self._env.step_dt + * self._env.reward_manager.get_term_cfg("tracking").weight + / (1 - self.cfg.gamma) + ) + soft = torch.softmax(-25.0 * self.bin_return / maximum_return, 0) + self.bin_weight.mul_(0.999).add_(0.001 * soft) + probabilities = 0.5 / motion.num_bins + 0.4 * soft + 0.1 * self.bin_weight / self.bin_weight.sum() + bins = torch.multinomial(probabilities, n, replacement=True) + start = motion.bin_start[bins] + end = torch.minimum(start + BIN_FRAMES, motion.clip_end[motion.bin_clip[bins]]) + frames = start + (torch.rand(n, device=self.device) * (end - start).float()).long() + + p, R, v, w, q, qd = (value.clone() for value in motion.ref_state(frames)) + if self.cfg.reset_noise and not self.cfg.evaluation: + p += torch.empty_like(p).uniform_(-1, 1) * p.new_tensor((0.05, 0.05, 0.01)) + rotation = torch.empty_like(p).uniform_(-1, 1) * p.new_tensor((0.1, 0.1, 0.2)) + R = so3_exp(rotation) @ R + v += torch.empty_like(v).uniform_(-1, 1) * v.new_tensor((0.5, 0.5, 0.2)) + w += torch.empty_like(w).uniform_(-1, 1) * w.new_tensor((0.52, 0.52, 0.78)) + q += torch.empty_like(q).uniform_(-0.1, 0.1) + q = q.clamp(self.kin.q_lo, self.kin.q_hi) + pose = torch.cat((p + self._env.scene.env_origins[ids], quat_from_matrix(R)), -1) + self.robot.write_root_link_pose_to_sim_index(root_pose=pose, env_ids=ids) + self.robot.write_root_link_velocity_to_sim_index(root_velocity=torch.cat((v, w), -1), env_ids=ids) + self.robot.write_joint_state_to_sim_index(position=q, velocity=qd, env_ids=ids) + self.robot.actuators.target_command.set_position_index(value=q, env_ids=ids) + self.start_frame[ids], self.start_bin[ids] = frames, bins + self._accounted_step[ids] = 0 + self.episode_return[ids] = 0 + self.discounted_return[ids] = 0 + if not self.cfg.evaluation: + self.xi[ids, 1:30] = torch.empty(n, 29, device=self.device).uniform_(-0.01, 0.01) + # The native environment clears episode_length_buf after manager resets. + self._command[ids] = motion.goal(frames, self.cfg.goal_frames) + + +class ContextJointPositionAction(JointPositionAction): + def __init__(self, cfg, env): + super().__init__(cfg, env) + self.motion = env.command_manager.get_term("motion") + + def process_actions(self, actions): + super().process_actions(actions) + self._processed_actions.add_(self.motion.xi[:, 1:30]) + + +def randomize_context(env, env_ids): + command = env.command_manager.get_term("motion") + if command.cfg.evaluation: + command.xi[:, 0] = 1.0 + else: + command.xi[:, 0].uniform_(0.3, 1.2) + command.xi[:, 30:33] = torch.empty(env.num_envs, 3, device=env.device).uniform_(-1, 1) * command.xi.new_tensor( + (0.025, 0.05, 0.05) + ) + default_com = command.robot.data.body_com_pos_b.torch[:, command.torso_id].clone() + command.robot.set_coms_index( + coms=(default_com + command.xi[:, 30:33]).unsqueeze(1), body_ids=[command.torso_id] + ) + model = command.physics.get_model() + shape = command.robot_view.get_attribute("shape_material_mu", model).shape + values = command.xi[:, 0].reshape(env.num_envs, 1, 1).expand(shape).contiguous() + command.robot_view.set_attribute("shape_material_mu", model, wp.from_torch(values)) + command.physics.add_model_change(ModelFlags.SHAPE_PROPERTIES) + + +def joint_pos_rel(env): + command = env.command_manager.get_term("motion") + return command.robot.data.joint_pos.torch - command.mdp.default_q - command.xi[:, 1:30] + + +def relative_orientation(env): + command = env.command_manager.get_term("motion") + R = matrix_from_quat(command.robot.data.root_link_pose_w.torch[:, 3:7]) + return mat_to_6d(R.transpose(-1, -2) @ command.motion.root_R[command.frame]) + + +def tracking_reward(env): + command = env.command_manager.get_term("motion") + return command.mdp.reward(command.state(), command.frame, env.action_manager.action, env.action_manager.prev_action) + + +def tracking_failure(env): + command = env.command_manager.get_term("motion") + return command.mdp.termination(command.state(), command.frame) + + +def clip_end(env): + command = env.command_manager.get_term("motion") + return command.frame >= command.motion.clip_end[command.motion.clip_id[command.start_frame]] + + +class TransitionRecorder(RecorderTerm): + """Publish actual pre/post tensors until the next step; optionally retain them as episodes. + + An exported episode's end is also a boundary, including explicit resets that happen after a + transition was retained. Online consumers additionally check adjacent state/step continuity. + """ + + def __init__(self, cfg, env): + super().__init__(cfg, env) + self.motion = env.command_manager.get_term("motion") + self.robot = self.motion.robot + self.latest = { + "s": torch.empty(self.num_envs, STATE_DIM, device=self.device), + "s1": torch.empty(self.num_envs, STATE_DIM, device=self.device), + "a": torch.empty(self.num_envs, 29, device=self.device), + "a_prev": torch.empty(self.num_envs, 29, device=self.device), + "xi": torch.empty_like(self.motion.xi), + "f": torch.empty(self.num_envs, dtype=torch.long, device=self.device), + "ep_step": torch.empty(self.num_envs, dtype=torch.int32, device=self.device), + "r": torch.empty(self.num_envs, device=self.device), + "done": torch.empty(self.num_envs, dtype=torch.bool, device=self.device), + "term": torch.empty(self.num_envs, dtype=torch.bool, device=self.device), + } + self._has_transition = False + + def _capture_state(self, target): + data = self.robot.data + target[:, :7].copy_(data.root_link_pose_w.torch) + target[:, :3].sub_(self._env.scene.env_origins) + target[:, 7:13].copy_(data.root_link_vel_w.torch) + target[:, 13:42].copy_(data.joint_pos.torch) + target[:, 42:71].copy_(data.joint_vel.torch) + + def record_pre_step(self): + self._capture_state(self.latest["s"]) + self.latest["a"].copy_(self._env.action_manager.action) + self.latest["a_prev"].copy_(self._env.action_manager.prev_action) + self.latest["xi"].copy_(self.motion.xi) + self.latest["f"].copy_(self.motion.frame) + self.latest["ep_step"].copy_(self._env.episode_length_buf) + self._has_transition = True + return None, None + + def record_post_step(self): + self._capture_state(self.latest["s1"]) + self.latest["r"].copy_(self._env.reward_buf) + self.latest["done"].copy_(self._env.reset_buf) + self.latest["term"].copy_(self._env.reset_terminated) + return ("transition", self.latest) if self.cfg.record_to_episodes else (None, None) + + def record_pre_reset(self, env_ids): + if self._has_transition: + self.latest["done"][slice(None) if env_ids is None else env_ids] = True + return None, None diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/motion.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/motion.py new file mode 100644 index 000000000000..eec9b12a6d95 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/motion.py @@ -0,0 +1,165 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""G1 reference motions resampled to 50 Hz, with bounded-memory FK construction.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F + +from isaaclab.utils.math import matrix_from_quat + +from .kinematics import G1Kinematics, mat_to_6d, so3_log, yaw_mat + +CONTROL_DT = 0.02 +BIN_FRAMES = 50 # 1 s bins for return-based reset sampling + + +def resample_motion(raw: torch.Tensor, src_fps: float): + """Resample metric/xyzw/radian G1 poses to the control rate using shortest-arc SLERP.""" + p, quat, j = raw[:, :3], raw[:, 3:7], raw[:, 7:36] + # resample to 50 Hz + T_src = raw.shape[0] + dur = (T_src - 1) / src_fps + n = int(np.floor(dur / CONTROL_DT)) + 1 + if n < 2: + raise ValueError("Motion is shorter than one control transition.") + t = torch.arange(n, dtype=torch.float32) * CONTROL_DT * src_fps + i0 = torch.clamp(t.floor().long(), max=T_src - 2) + a = (t - i0.float()).clamp(0, 1) + p = p[i0] * (1 - a[:, None]) + p[i0 + 1] * a[:, None] + j = j[i0] * (1 - a[:, None]) + j[i0 + 1] * a[:, None] + # Batched shortest-arc SLERP; the native helper accepts a single interpolation pair. + q0, q1 = quat[i0], quat[i0 + 1] + dot = (q0 * q1).sum(-1, keepdim=True) + q1 = torch.where(dot < 0, -q1, q1) + theta = torch.acos(dot.abs().clamp(max=1.0)) + sine, weight = torch.sin(theta), a[:, None] + w0 = torch.where(sine < 1e-6, 1 - weight, torch.sin((1 - weight) * theta) / (sine + 1e-12)) + w1 = torch.where(sine < 1e-6, weight, torch.sin(weight * theta) / (sine + 1e-12)) + quat = F.normalize(w0 * q0 + w1 * q1, dim=-1) + return p, quat, j + + +class MotionLib: + def __init__(self, path: str, kin: G1Kinematics, device: str, src_fps: float = 30.0): + if not np.isfinite(src_fps) or src_fps <= 0: + raise ValueError("Source frame rate must be finite and positive.") + files = sorted((*Path(path).rglob("*.csv"), *Path(path).rglob("*.npy"))) + if not files: + raise ValueError(f"No motion CSV/NumPy files found in {path}.") + self.device = device + self.names = [f.relative_to(path).with_suffix("").as_posix() for f in files] + if len(set(self.names)) != len(self.names): + raise ValueError("Motion names must be unique across CSV and NumPy files.") + ps, qs, js = [], [], [] + lens = [] + for f in files: + raw = torch.tensor( + np.load(f, allow_pickle=False) if f.suffix == ".npy" else np.loadtxt(f, delimiter=","), + dtype=torch.float32, + ) + if raw.ndim != 2 or raw.shape[0] < 2 or raw.shape[1] != 36: + raise ValueError(f"Motion {f} must contain at least two 36-channel frames.") + if not torch.isfinite(raw).all() or (raw[:, 3:7].norm(dim=-1) < 1e-6).any(): + raise ValueError(f"Motion {f} contains nonfinite values or invalid quaternions.") + if src_fps == 1.0 / CONTROL_DT: + p, quat, j = raw[:, :3], F.normalize(raw[:, 3:7], dim=-1), raw[:, 7:] + else: + p, quat, j = resample_motion(raw, src_fps) + n = len(p) + ps.append(p), qs.append(quat), js.append(j) + lens.append(n) + self.clip_len = torch.tensor(lens, device=device) + self.clip_start = torch.cumsum( + torch.cat((torch.zeros(1, dtype=torch.long), torch.tensor(lens[:-1], dtype=torch.long))), 0 + ).to(device) # dtype: one clip gives an empty (float) tensor + self.clip_end = self.clip_start + self.clip_len - 1 # last valid frame (inclusive) + self.num_frames = int(sum(lens)) + self.clip_id = torch.cat([torch.full((n,), i, dtype=torch.long) for i, n in enumerate(lens)]).to(device) + + root_p, root_R, root_v, root_w, jq, jqd = [], [], [], [], [], [] + for p, quat, j in zip(ps, qs, js): + R = matrix_from_quat(quat) + v = torch.gradient(p, spacing=CONTROL_DT, dim=0)[0] + qd = torch.gradient(j, spacing=CONTROL_DT, dim=0)[0] + # angular velocity (world) by central differences of rotations + w = torch.zeros_like(p) + w[1:-1] = so3_log(R[2:] @ R[:-2].transpose(-1, -2)) / (2 * CONTROL_DT) + w[0] = so3_log(R[1] @ R[0].T) / CONTROL_DT + w[-1] = so3_log(R[-1] @ R[-2].T) / CONTROL_DT + root_p.append(p), root_R.append(R), root_v.append(v), root_w.append(w), jq.append(j), jqd.append(qd) + cat = lambda xs: torch.cat(xs).to(device) # noqa: E731 + self.root_p, self.root_R, self.root_v, self.root_w = cat(root_p), cat(root_R), cat(root_v), cat(root_w) + self.q, self.qd = cat(jq), cat(jqd) + with torch.no_grad(): + for s in range(0, self.num_frames, 65536): + e = min(s + 65536, self.num_frames) + out = kin.forward( + self.root_p[s:e], self.root_R[s:e], self.root_v[s:e], self.root_w[s:e], self.q[s:e], self.qd[s:e] + ) + if s == 0: + self.body_p, self.body_R, self.body_v, self.body_w = ( + value.new_empty((self.num_frames, *value.shape[1:])) for value in out + ) + for destination, value in zip((self.body_p, self.body_R, self.body_v, self.body_w), out, strict=True): + destination[s:e].copy_(value) + # per-frame static goal parts: q, qd, root velocities in the root frame + Rt = self.root_R.transpose(-1, -2) + self.root_v_local = (Rt @ self.root_v.unsqueeze(-1)).squeeze(-1) + self.root_w_local = (Rt @ self.root_w.unsqueeze(-1)).squeeze(-1) + self.root_gz = self.root_R[:, 2, 2] # -(projected gravity z) of the reference root: g_b = R^T (0,0,-1) + # 1 s bins + bin_clip, bin_start, frame_bins = [], [], [] + frame_offset = 0 + for c in range(len(lens)): + nb = (lens[c] - 2) // BIN_FRAMES + 1 + frame_bins.append((torch.arange(lens[c]) // BIN_FRAMES).clamp_max(nb - 1) + len(bin_start)) + for b in range(nb): + bin_clip.append(c) + bin_start.append(frame_offset + b * BIN_FRAMES) + frame_offset += lens[c] + self.bin_clip = torch.tensor(bin_clip, device=device) + self.bin_start = torch.tensor(bin_start, device=device) + self.num_bins = len(bin_start) + self.frame_bin = torch.cat(frame_bins).to(device) + + def future_frames(self, f: torch.Tensor, K: int) -> torch.Tensor: + """Frames f+1..f+K clamped to the clip's last frame, shape (B, K).""" + end = self.clip_end[self.clip_id[f]] + fut = f.unsqueeze(-1) + torch.arange(1, K + 1, device=f.device) + return torch.minimum(fut, end.unsqueeze(-1)) + + def goal(self, f: torch.Tensor, K: int) -> torch.Tensor: + """Goal features from the reference only: K future frames relative to the current reference frame.""" + fut = self.future_frames(f, K) # (B,K) + # Beyond the recorded clip, the padded pose is held still rather than continuing at its endpoint velocity. + recorded = (fut == f.unsqueeze(-1) + torch.arange(1, K + 1, device=f.device)).unsqueeze(-1) + H = yaw_mat(self.root_R[f]) # (B,3,3) + Ht = H.transpose(-1, -2).unsqueeze(1) + dp = (Ht @ (self.root_p[fut] - self.root_p[f].unsqueeze(1)).unsqueeze(-1)).squeeze(-1) + rot = mat_to_6d(Ht @ self.root_R[fut]) + feats = torch.cat( + ( + dp, + rot, + self.root_v_local[fut] * recorded, + self.root_w_local[fut] * recorded, + self.q[fut], + self.qd[fut] * (0.1 * recorded), + ), + -1, + ) # (B,K,73) + return feats.reshape(f.shape[0], -1) + + def ref_state(self, f: torch.Tensor): + return self.root_p[f], self.root_R[f], self.root_v[f], self.root_w[f], self.q[f], self.qd[f] diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/tracking.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/tracking.py new file mode 100644 index 000000000000..859fd256b95d --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/tracking.py @@ -0,0 +1,128 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""Analytic MDP heads shared by the simulator and the world model (paper Appendix A). + +State tuple ``S = (p, R, v, w, q, qd)``: pelvis position, rotation matrix, world linear/angular +velocity of the pelvis link frame, joint positions and velocities. +""" + +from __future__ import annotations + +import torch + +from isaaclab.utils import configclass +from isaaclab.utils.math import matrix_from_quat + +from .kinematics import G1Kinematics, mat_to_6d, rot_angle_sq, yaw_mat +from .motion import MotionLib + +STATE_DIM = 71 +OBS_DIM = 29 + 29 + 3 + 3 + 29 + 6 +XI_DIM = 33 +NUM_J = 29 + +MAX_REWARD_RATE = 5.0 + + +@configclass +class RewardCfg: + body_pos_std: float = 0.3 + body_ori_std: float = 0.4 + body_linvel_std: float = 1.0 + body_angvel_std: float = 3.14 + root_pos_std: float = 0.3 + root_ori_std: float = 0.4 + w_action_rate: float = -0.1 + w_joint_limit: float = -2.0 + soft_limit_factor: float = 0.9 + term_height: float = 0.25 + term_tilt: float = 0.8 + + +def unflatten_state(x: torch.Tensor): + return x[..., 0:3], matrix_from_quat(x[..., 3:7]), x[..., 7:10], x[..., 10:13], x[..., 13:42], x[..., 42:71] + + +class TrackingObjective: + def __init__(self, kin: G1Kinematics, motion: MotionLib, default_q: torch.Tensor, cfg: RewardCfg | None = None): + self.kin, self.motion, self.cfg = kin, motion, cfg if cfg is not None else RewardCfg() + cfg = self.cfg + self.default_q = default_q + mid, half = (kin.q_lo + kin.q_hi) / 2, (kin.q_hi - kin.q_lo) / 2 + self.soft_lo, self.soft_hi = mid - half * cfg.soft_limit_factor, mid + half * cfg.soft_limit_factor + + # ---------------- observation ---------------- + def obs(self, S, a_prev, xi, f) -> torch.Tensor: + """Noise-free proprioceptive observation (+ reference-relative root orientation).""" + p, R, v, w, q, qd = S + Rt = R.transpose(-1, -2) + dq0 = xi[..., 1:30] + g_b = -R[..., 2, :] # R^T (0,0,-1) + rel = mat_to_6d(Rt @ self.motion.root_R[f]) + return torch.cat((q - self.default_q - dq0, qd, (Rt @ w.unsqueeze(-1)).squeeze(-1), g_b, a_prev, rel), -1) + + # ---------------- reward / termination ---------------- + def reward(self, S, f, a, a_prev, terms: bool = False): + """Reward rate comparing S with reference frame f; the caller applies its step duration.""" + c, m = self.cfg, self.motion + p, R, v, w, q, qd = S + bp, bR, bv, bw = self.kin.forward(p, R, v, w, q, qd) + rp, rR = m.root_p[f], m.root_R[f] + # reference bodies re-anchored to the robot's root (xy + yaw), as in BeyondMimic + dR = yaw_mat(R @ rR.transpose(-1, -2)) + anchor = torch.cat((p[..., :2], rp[..., 2:3]), -1) + ref_bp = anchor.unsqueeze(-2) + ( + (dR.unsqueeze(-3) @ (m.body_p[f] - rp.unsqueeze(-2)).unsqueeze(-1)).squeeze(-1) + ) + ref_bR = dR.unsqueeze(-3) @ m.body_R[f] + e_pos = ((ref_bp - bp) ** 2).sum(-1).mean(-1) + e_ori = rot_angle_sq(ref_bR, bR).mean(-1) + e_lv = ((m.body_v[f] - bv) ** 2).sum(-1).mean(-1) + e_av = ((m.body_w[f] - bw) ** 2).sum(-1).mean(-1) + e_rp = ((rp - p) ** 2).sum(-1) + e_ro = rot_angle_sq(rR, R) + r_track = ( + torch.exp(-e_pos / c.body_pos_std**2) + + torch.exp(-e_ori / c.body_ori_std**2) + + torch.exp(-e_lv / c.body_linvel_std**2) + + torch.exp(-e_av / c.body_angvel_std**2) + + 0.5 * torch.exp(-e_rp / c.root_pos_std**2) + + 0.5 * torch.exp(-e_ro / c.root_ori_std**2) + ) + r_rate = ((a - a_prev) ** 2).sum(-1) + r_lim = (torch.relu(q - self.soft_hi) + torch.relu(self.soft_lo - q)).sum(-1) + r = r_track + c.w_action_rate * r_rate + c.w_joint_limit * r_lim + if terms: + return r, { + "track": r_track, + "e_body_pos": e_pos, + "e_body_ori": e_ori, + "e_root_pos": e_rp, + "action_rate": r_rate, + "joint_limit": r_lim, + } + return r + + def termination(self, S, f) -> torch.Tensor: + p, R = S[0], S[1] + rp, rR = self.motion.root_p[f], self.motion.root_R[f] + e_z = (rp[..., 2] - p[..., 2]).abs() + e_tilt = (rR[..., 2, 2] - R[..., 2, 2]).abs() + return (e_z > self.cfg.term_height) | (e_tilt > self.cfg.term_tilt) + + def soft_termination(self, S, f, tz: float = 0.02, tt: float = 0.04) -> torch.Tensor: + """Differentiable termination probability for imagined rollouts.""" + p, R = S[0], S[1] + rp, rR = self.motion.root_p[f], self.motion.root_R[f] + e_z = (rp[..., 2] - p[..., 2]).abs() + e_tilt = (rR[..., 2, 2] - R[..., 2, 2]).abs() + alive = (1 - torch.sigmoid((e_z - self.cfg.term_height) / tz)) * ( + 1 - torch.sigmoid((e_tilt - self.cfg.term_tilt) / tt) + ) + return 1 - alive diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/digit/rough_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/digit/rough_env_cfg.py index 58b2ff1cad2f..cbd1e3865774 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/digit/rough_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/digit/rough_env_cfg.py @@ -40,12 +40,8 @@ class DigitPhysicsCfg(PresetCfg): gpu_total_aggregate_pairs_capacity=2**23, ) physx = PhysxAutoCfg(isaacsim_physx=isaacsim_physx) - # ``class_type`` is reset to ``None`` because ``NewtonCfg.__post_init__`` re-derives it from - # ``solver_cfg`` and refuses an explicit value -- ``replace()`` would otherwise carry the - # already-derived value from the source instance forward as one. newton_mjwarp = replace( _ROUGH_NEWTON_MJWARP, - class_type=None, solver_cfg=replace(_ROUGH_NEWTON_MJWARP.solver_cfg, njmax=5000, nconmax=2000), ) default = isaacsim_physx diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/__init__.py index 99b7549a6463..b70eb314c951 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/__init__.py @@ -14,6 +14,23 @@ from . import agents +for suffix, config in ( + ("", "CartpoleNeRDEnvCfg"), + ("-Explore", "CartpoleNeRDExplorationEnvCfg"), + ("-Learned", "CartpoleNeRDLearnedEnvCfg"), +): + gym.register( + id=f"Isaac-Cartpole-NeRD{suffix}", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.cartpole_nerd_env_cfg:{config}", + "nerd_cfg_entry_point": "isaaclab_rl.nerd:NeRDTrainCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CartpolePPORunnerCfg", + "default_agent": "rsl_rl", + }, + ) + ## # Register Gym environments -- direct workflow. ## diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/cartpole_nerd_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/cartpole_nerd_env_cfg.py new file mode 100644 index 000000000000..5572cb0d9820 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/cartpole_nerd_env_cfg.py @@ -0,0 +1,65 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Native Cartpole tasks for offline dynamics collection and learned-physics policy training.""" + +import math + +from isaaclab_newton.envs.mdp.recorders import NeRDRecorderCfg, NeRDRecorderManagerCfg +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +from isaaclab_newton.physics.nerd_manager_cfg import NeRDSolverCfg + +from isaaclab.actuators import IdealPDActuatorCfg +from isaaclab.utils import configclass + +from .cartpole_manager_env_cfg import CartpoleEnvCfg + + +@configclass +class CartpoleNeRDEnvCfg(CartpoleEnvCfg): + """MuJoCo Warp teacher at 120 Hz, with optional native transition capture.""" + + recorders: NeRDRecorderManagerCfg = NeRDRecorderManagerCfg() + + def __post_init__(self): + super().__post_init__() + self.scene.robot.actuators = { + "joints": IdealPDActuatorCfg( + joint_names_expr=[".*"], + stiffness=0.0, + damping={"slider_to_cart": 10.0, "cart_to_pole": 0.0}, + actuator_effort_limit=400.0, + ) + } + self.decimation = self.sim.render_interval = 1 + self.sim.dt = 1.0 / 120.0 + self.sim.physics = NewtonCfg( + solver_cfg=MJWarpSolverCfg(njmax=5, nconmax=3, cone="pyramidal", impratio=1.0, integrator="implicitfast"), + num_substeps=1, + use_cuda_graph=True, + ) + + +@configclass +class CartpoleNeRDExplorationEnvCfg(CartpoleNeRDEnvCfg): + """Task-independent initial states for a dynamics dataset, including full pole rotations.""" + + def __post_init__(self): + super().__post_init__() + self.episode_length_s = 2.0 + self.events.reset_cart_position.params.update(position_range=(-2.0, 2.0), velocity_range=(-2.0, 2.0)) + self.events.reset_pole_position.params.update( + position_range=(-math.pi, math.pi), velocity_range=(-2 * math.pi, 2 * math.pi) + ) + self.recorders.transition = NeRDRecorderCfg() + + +@configclass +class CartpoleNeRDLearnedEnvCfg(CartpoleNeRDEnvCfg): + """The same control task advanced by a NeRD checkpoint.""" + + def __post_init__(self): + super().__post_init__() + self.sim.physics = NewtonCfg(solver_cfg=NeRDSolverCfg(), num_substeps=1, use_cuda_graph=True) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/anymal_c/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/anymal_c/__init__.py new file mode 100644 index 000000000000..aa2b0b08095d --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/anymal_c/__init__.py @@ -0,0 +1,41 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Native ANYmal C NeRD collection and velocity-tracking tasks.""" + +import gymnasium as gym + +for suffix, config in ( + ("", "AnymalCNeRDEnvCfg"), + ("-Explore", "AnymalCNeRDExplorationEnvCfg"), + ("-Learned", "AnymalCNeRDLearnedEnvCfg"), +): + gym.register( + id=f"Isaac-Velocity-Flat-Anymal-C-NeRD{suffix}", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.nerd_env_cfg:{config}", + "nerd_cfg_entry_point": f"{__name__}:nerd_cfg.yaml", + "rsl_rl_cfg_entry_point": "isaaclab_tasks.core.velocity.config.anymal_d.agents.rsl_rl_ppo_cfg:AnymalDFlatPPORunnerCfg", + "default_agent": "rsl_rl", + }, + ) + +for suffix, config in ( + ("", "AnymalCNeRDRoughEnvCfg"), + ("-Explore", "AnymalCNeRDRoughExplorationEnvCfg"), + ("-Learned", "AnymalCNeRDRoughLearnedEnvCfg"), +): + gym.register( + id=f"Isaac-Velocity-Rough-Anymal-C-NeRD{suffix}", + entry_point="isaaclab.envs:ManagerBasedRLEnv", disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.nerd_env_cfg:{config}", + "nerd_cfg_entry_point": f"{__name__}.nerd_env_cfg:AnymalCRoughNeRDTrainCfg", + "rsl_rl_cfg_entry_point": "isaaclab_tasks.core.velocity.config.anymal_d.agents.rsl_rl_ppo_cfg:AnymalDFlatPPORunnerCfg", + "default_agent": "rsl_rl", + }, + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/anymal_c/nerd_cfg.yaml b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/anymal_c/nerd_cfg.yaml new file mode 100644 index 000000000000..dfcd176516ba --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/anymal_c/nerd_cfg.yaml @@ -0,0 +1,17 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +max_updates: 60000 +learning_rate_decay_updates: 60000 +batch_size: 4096 +validation_fraction: 0.05 +network: + num_heads: 10 + head_dim: 64 + intermediate_size: 896 + head_hidden: 192 + temporal_attention_spans: [4, 4, 4, 4] + rms_norm_eps: 1.0e-6 + linear_probe: true diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/anymal_c/nerd_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/anymal_c/nerd_env_cfg.py new file mode 100644 index 000000000000..f95f1cecd3fe --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/anymal_c/nerd_env_cfg.py @@ -0,0 +1,212 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""ANYmal C applied-effort dynamics and locomotion tasks.""" + +import math + +import torch +from isaaclab_newton.envs.mdp.recorders import NeRDRecorderCfg, NeRDRecorderManagerCfg +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg, NewtonShapeCfg +from isaaclab_newton.physics.nerd_manager_cfg import NeRDSolverCfg + +import isaaclab.envs.mdp as mdp +import isaaclab.sim as sim_utils +from isaaclab.assets import AssetBaseCfg +from isaaclab.managers import ManagerTermBase, RewardTermCfg, TerminationTermCfg +from isaaclab.terrains import FlatPatchSamplingCfg, TerrainImporterCfg +from isaaclab.terrains.config.rough import ROUGH_TERRAINS_CFG +from isaaclab.utils import configclass + +from isaaclab_assets.robots.anymal import ANYMAL_C_CFG + +from ..anymal_d.flat_env_cfg import AnymalDFlatEnvCfg + + +class control_period_joint_acc_l2(ManagerTermBase): + """Difference velocities over the control period in both numerical and learned physics.""" + + def __init__(self, cfg, env): + super().__init__(cfg, env) + self.robot = env.scene["robot"] + self.previous = self.robot.data.joint_vel.torch.clone() + self.fresh = torch.ones(env.num_envs, dtype=torch.bool, device=env.device) + + def reset(self, env_ids=None): + self.fresh[slice(None) if env_ids is None else env_ids] = True + + def __call__(self, env): + velocity = self.robot.data.joint_vel.torch + penalty = torch.where(self.fresh, 0.0, ((velocity - self.previous) / env.step_dt).square().sum(dim=1)) + self.previous.copy_(velocity) + self.fresh.zero_() + return penalty + + +@configclass +class AnymalCNeRDEnvCfg(AnymalDFlatEnvCfg): + """Flat-ground teacher with 5 ms ANYdrive updates and a 20 ms policy control period. + + The policy objective uses state-derived quantities that the learned solver also produces. + Fixed probe geometry requires one plane per world. Rough terrain needs another feature + contract and is not represented by these flat-ground probes. + """ + + recorders: NeRDRecorderManagerCfg = NeRDRecorderManagerCfg() + + def __post_init__(self): + super().__post_init__() + self.scene.robot = ANYMAL_C_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + self.sim.physics = NewtonCfg( + solver_cfg=MJWarpSolverCfg( + njmax=60, nconmax=25, cone="elliptic", impratio=100.0, integrator="implicitfast" + ), + num_substeps=1, + use_cuda_graph=True, + default_shape_cfg=NewtonShapeCfg(margin=0.0, ke=160000.0, kd=1100.0), + ) + material = self.scene.terrain.physics_material + self.scene.terrain = None + self.scene.ground = AssetBaseCfg( + prim_path="{ENV_REGEX_NS}/Ground", + spawn=sim_utils.GroundPlaneCfg( + size=(self.scene.env_spacing, self.scene.env_spacing), physics_material=material + ), + ) + self.scene.contact_forces = None + self.rewards.feet_air_time = self.rewards.undesired_contacts = self.rewards.dof_torques_l2 = None + self.rewards.dof_acc_l2 = RewardTermCfg(func=control_period_joint_acc_l2, weight=-1.2e-6) + self.terminations.base_contact = None + self.terminations.bad_orientation = TerminationTermCfg( + func=mdp.bad_orientation, params={"limit_angle": math.pi / 2} + ) + self.events.physics_material = self.events.add_base_mass = self.events.base_com = None + self.events.base_external_force_torque = self.events.push_robot = None + + +@configclass +class AnymalCNeRDExplorationEnvCfg(AnymalCNeRDEnvCfg): + """Random poses and targets, retaining falls for task-independent dynamics collection.""" + + def __post_init__(self): + super().__post_init__() + self.decimation = 1 + self.episode_length_s = 2.0 + self.terminations.bad_orientation = None + self.events.reset_base.params = { + "pose_range": { + "x": (-0.5, 0.5), + "y": (-0.5, 0.5), + "z": (-0.1, 0.1), + "roll": (-0.3, 0.3), + "pitch": (-0.3, 0.3), + "yaw": (-math.pi, math.pi), + }, + "velocity_range": { + "x": (-1.0, 1.0), + "y": (-1.0, 1.0), + "z": (-0.5, 0.5), + "roll": (-1.0, 1.0), + "pitch": (-1.0, 1.0), + "yaw": (-1.0, 1.0), + }, + } + self.events.reset_robot_joints.func = mdp.reset_joints_by_offset + self.events.reset_robot_joints.params = {"position_range": (-0.5, 0.5), "velocity_range": (-1.0, 1.0)} + self.recorders.transition = NeRDRecorderCfg( + heading=True, + contact_bodies=("LF_FOOT", "LH_FOOT", "RF_FOOT", "RH_FOOT"), + root_body="base", + ) + + +@configclass +class AnymalCNeRDLearnedEnvCfg(AnymalCNeRDEnvCfg): + """Keep ANYdrive and run four 5 ms learned physics steps per 20 ms policy action.""" + + def __post_init__(self): + super().__post_init__() + self.sim.physics = self.sim.physics.replace(solver_cfg=NeRDSolverCfg()) + + +@configclass +class AnymalCNeRDRoughExplorationEnvCfg(AnymalCNeRDExplorationEnvCfg): + """Physics-step collection across rough surfaces, conditioned on all Newton rigid contacts.""" + + def __post_init__(self): + super().__post_init__() + material = self.scene.ground.spawn.physics_material + self.scene.ground = None + self.scene.terrain = TerrainImporterCfg( + prim_path="/World/ground", + terrain_type="generator", + max_init_terrain_level=None, + terrain_generator=ROUGH_TERRAINS_CFG.replace( + seed=42, num_rows=4, num_cols=6, border_width=2.0, curriculum=False + ), + physics_material=material, + ) + self.sim.physics.solver_cfg.use_mujoco_contacts = False + self.sim.physics.solver_cfg.nconmax = 64 + self.sim.physics.solver_cfg.njmax = 256 + for terrain in self.scene.terrain.terrain_generator.sub_terrains.values(): + terrain.flat_patch_sampling = { + "init_pos": FlatPatchSamplingCfg( + num_patches=64, patch_radius=0.5, max_height_diff=0.12, x_range=(-3.0, 3.0), y_range=(-3.0, 3.0) + ) + } + self.scene.robot.init_state.pos = (0.0, 0.0, 0.8) + self.events.reset_base.func = mdp.reset_root_state_from_terrain + self.events.reset_base.params["pose_range"] = { + "roll": (-0.3, 0.3), + "pitch": (-0.3, 0.3), + "yaw": (-math.pi, math.pi), + } + self.recorders.transition.contact_bodies = () + self.recorders.transition.max_contacts = 64 + + +@configclass +class AnymalCNeRDRoughEnvCfg(AnymalCNeRDRoughExplorationEnvCfg): + """Rough-terrain locomotion with native ANYdrive and proprioceptive policy observations.""" + + def __post_init__(self): + super().__post_init__() + self.recorders.transition = None + self.decimation = self.sim.render_interval = 4 + self.episode_length_s = 20.0 + self.terminations.bad_orientation = TerminationTermCfg( + func=mdp.bad_orientation, params={"limit_angle": math.pi / 2} + ) + + +@configclass +class AnymalCNeRDRoughLearnedEnvCfg(AnymalCNeRDRoughEnvCfg): + """Replace the numerical rough-terrain solver while retaining collisions and ANYdrive.""" + + def __post_init__(self): + super().__post_init__() + self.sim.physics = self.sim.physics.replace(solver_cfg=NeRDSolverCfg()) + + +@configclass +class AnymalCRoughNeRDTrainCfg: + """Contact-conditioned offline fitting; optimizer defaults belong to the shared learner.""" + + max_updates = 60000 + learning_rate_decay_updates = 60000 + batch_size = 4096 + precision = "bfloat16" + network = { + "num_heads": 10, + "head_dim": 64, + "intermediate_size": 896, + "head_hidden": 192, + "temporal_attention_spans": (4, 4, 4, 4), + "rms_norm_eps": 1.0e-6, + "linear_probe": True, + "contact_slot_count": 8, + "contact_slot_dim": 64, + } diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/velocity_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/velocity_env_cfg.py index 0e92889fc9fd..573d6cdc8c12 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/velocity_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/velocity_env_cfg.py @@ -377,13 +377,14 @@ def __post_init__(self): def play_mode(self): """Play-mode overrides shared by the velocity-tracking environments.""" super().play_mode() - # spawn the robot randomly in the grid (instead of their terrain levels) - self.scene.terrain.max_init_terrain_level = None - # reduce the number of terrains to save memory - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False + if self.scene.terrain is not None: + # spawn the robot randomly in the grid (instead of their terrain levels) + self.scene.terrain.max_init_terrain_level = None + # reduce the number of terrains to save memory + if self.scene.terrain.terrain_generator is not None: + self.scene.terrain.terrain_generator.num_rows = 5 + self.scene.terrain.terrain_generator.num_cols = 5 + self.scene.terrain.terrain_generator.curriculum = False # remove random pushing events self.events.base_external_force_torque = None self.events.push_robot = None diff --git a/source/isaaclab_tasks/test/test_g1_tracking.py b/source/isaaclab_tasks/test/test_g1_tracking.py new file mode 100644 index 000000000000..2d9072c5e254 --- /dev/null +++ b/source/isaaclab_tasks/test/test_g1_tracking.py @@ -0,0 +1,212 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""G1 capture and motion contracts, including stationary padding beyond recorded clip ends.""" + +import os +from pathlib import Path +from types import SimpleNamespace + +import gymnasium as gym +import h5py +import mujoco +import numpy as np +import pytest +import torch +import warp as wp + +from isaaclab.app import launch_simulation +from isaaclab.managers import DatasetExportMode, RewardTermCfg, TerminationTermCfg +from isaaclab.test.utils import DeviceScope, test_devices +from isaaclab.utils.math import matrix_from_quat +from isaaclab.utils.string import resolve_matching_names_values + +from isaaclab_tasks.contrib.g1_tracking.g1_tracking_env_cfg import TransitionRecorderCfg +from isaaclab_tasks.contrib.g1_tracking.kinematics import TRACKED_BODIES +from isaaclab_tasks.contrib.g1_tracking.motion import MotionLib +from isaaclab_tasks.utils import resolve_task_config + + +def test_motion_rates_and_stationary_endpoint_padding(tmp_path): + raw = np.zeros((121, 36), dtype=np.float32) + raw[:, 2], raw[:, 6] = 0.8, 1.0 + raw[:, 0] = np.linspace(0, 0.2, len(raw)) + roll = np.linspace(0, 0.2, len(raw)) + raw[:, 3], raw[:, 6] = np.sin(roll / 2), np.cos(roll / 2) + raw[:, 7] = np.linspace(0, 0.1, len(raw)) + for directory in ("a", "b"): + (tmp_path / directory).mkdir() + np.savetxt(tmp_path / directory / "walk.csv", raw, delimiter=",") + # Kinematics is an independent dependency here; this test checks temporal pose/rate semantics. + kin = SimpleNamespace(forward=lambda p, R, v, w, q, qd: (p[:, None], R[:, None], v[:, None], w[:, None])) + motion = MotionLib(str(tmp_path), kin, "cpu", src_fps=120) + assert motion.names == ["a/walk", "b/walk"] + assert motion.clip_len.tolist() == [51, 51] and motion.num_bins == 2 + torch.testing.assert_close(motion.root_v, torch.tensor([0.2, 0, 0]).expand(102, -1), atol=2e-6, rtol=1e-5) + assert motion.future_frames(torch.tensor([49, 100]), 3).tolist() == [[50, 50, 50], [101, 101, 101]] + goals = motion.goal(torch.tensor([47, 49, 50, 100]), 3).reshape(4, 3, 73) + # Recorded endpoint rates remain factual; only slots beyond the clip describe a stationary held pose. + expected_velocity = torch.tensor([[0.2, 0.2, 0.2], [0.2, 0.0, 0.0], [0.0, 0.0, 0.0], [0.2, 0.0, 0.0]]) + torch.testing.assert_close(goals[:, :, 9], expected_velocity, atol=2e-6, rtol=1e-5) + torch.testing.assert_close(goals[:, :, 12], expected_velocity, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(goals[:, :, 44], expected_velocity * 0.05, atol=2e-6, rtol=1e-5) + torch.testing.assert_close(goals[1:, :, 15], torch.full((3, 3), 0.1), atol=2e-6, rtol=1e-5) + with pytest.raises(ValueError, match="frame rate"): + MotionLib(str(tmp_path), kin, "cpu", src_fps=0) + + +def never(env): + return torch.zeros(env.num_envs, dtype=torch.bool, device=env.device) + + +@pytest.fixture(scope="module", params=test_devices(DeviceScope.CUDA)) +def task(tmp_path_factory, request): + torch.backends.cuda.matmul.allow_tf32 = False + output = tmp_path_factory.mktemp("native_task") + cfg, _ = resolve_task_config( + "IsaacContrib-Motion-Tracking-G1", None, overrides=["env.scene.num_envs=40", "env.seed=123"] + ) + cfg.sim.device = request.param + data_dir = os.environ.get("ISAACLAB_G1_TRACKING_DATA") + if not data_dir: + pytest.skip("Set ISAACLAB_G1_TRACKING_DATA to the prepared G1 tracking directory.") + cfg.scene.robot.spawn.usd_path = str(Path(data_dir) / "g1_usd/g1_29dof/g1_29dof.usda") + cfg.commands.motion.kinematics_path = str(Path(data_dir) / "g1_description/g1_29dof.xml") + cfg.commands.motion.motion_dir = str(Path(data_dir) / "lafan1_g1") + cfg.recorders.transition = TransitionRecorderCfg() + cfg.observations.policy.enable_corruption = False + cfg.observations.policy.history_length = 3 + cfg.commands.motion.goal_frames = 2 + cfg.commands.motion.reset_noise = False + cfg.recorders.transition.record_to_episodes = True + cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_ALL + cfg.recorders.dataset_export_dir_path = str(output) + cfg.recorders.buffer_initial_capacity = 2 + for name in ("failed", "clip_end", "time_out"): + getattr(cfg.terminations, name).func = never + with launch_simulation(cfg): + env = gym.make("IsaacContrib-Motion-Tracking-G1", cfg=cfg).unwrapped + try: + yield env, output + finally: + env.close() + + +def test_real_capture_observations_context_and_offline_export(task): + env, output = task + command = env.command_manager.get_term("motion") + recorder = env.recorder_manager.get_term("transition") + observations, _ = env.reset() + assert env.cfg.sim.physics.use_cuda_graph + robot = command.robot + model = mujoco.MjModel.from_xml_path(command.cfg.kinematics_path) + expected_com = torch.tensor(model.body("torso_link").ipos, device=env.device, dtype=torch.float32) + torch.testing.assert_close( + robot.data.body_com_pos_b.torch[:, command.torso_id], expected_com + command.xi[:, 30:33] + ) + mu = wp.to_torch(command.robot_view.get_attribute("shape_material_mu", command.physics.get_model())) + torch.testing.assert_close(mu, command.xi[:, :1, None].expand_as(mu)) + previous_action = torch.zeros(env.num_envs, 29, device=env.device) + records = [] + for step in range(3): + state = torch.cat( + ( + robot.data.root_link_pose_w.torch, + robot.data.root_link_vel_w.torch, + robot.data.joint_pos.torch, + robot.data.joint_vel.torch, + ), + -1, + ).clone() + frame = command.frame.clone() + action = torch.full_like(previous_action, 0.01 * (step + 1)) + observations, reward, terminated, truncated, _ = env.step(action) + transition = recorder.latest + assert not (terminated | truncated).any() + torch.testing.assert_close(transition["s"], state, rtol=0, atol=0) + torch.testing.assert_close(transition["a_prev"], previous_action, rtol=0, atol=0) + torch.testing.assert_close(transition["a"], action, rtol=0, atol=0) + torch.testing.assert_close(transition["r"], reward, rtol=0, atol=0) + torch.testing.assert_close(transition["f"], frame, rtol=0, atol=0) + analytical = command.mdp.obs(command.state(), action, command.xi, command.frame) + torch.testing.assert_close(observations["policy"][:, -1], analytical, rtol=1e-5, atol=2e-6) + torch.testing.assert_close(observations["goal"], command.motion.goal(command.frame, 2)) + action_term = env.action_manager.get_term("joint_pos") + _, _, values = resolve_matching_names_values(env.cfg.actions.joint_pos.scale, robot.joint_names) + scale = torch.tensor(values, device=env.device) + expected_target = command.mdp.default_q + command.xi[:, 1:30] + scale * action + torch.testing.assert_close(action_term.processed_actions, expected_target) + records.append({key: value.clone() for key, value in transition.items()}) + previous_action = action + body_ids = [robot.body_names.index(name) for name in TRACKED_BODIES] + positions, rotations, velocities, angular = command.kin.forward(*command.state()) + torch.testing.assert_close(positions, robot.data.body_link_pos_w.torch[:, body_ids], atol=5e-5, rtol=1e-5) + torch.testing.assert_close( + rotations, matrix_from_quat(robot.data.body_link_quat_w.torch[:, body_ids]), atol=5e-5, rtol=1e-5 + ) + torch.testing.assert_close(velocities, robot.data.body_link_lin_vel_w.torch[:, body_ids], atol=1e-4, rtol=1e-4) + torch.testing.assert_close(angular, robot.data.body_link_ang_vel_w.torch[:, body_ids], atol=1e-4, rtol=1e-4) + snapshot = env.recorder_manager.get_episode(0).data["transition"] + for key in records[0]: + torch.testing.assert_close(snapshot[key], torch.stack([record[key][0] for record in records]), rtol=0, atol=0) + env.recorder_manager.export_episodes([0]) + with h5py.File(output / "dataset.hdf5", "r") as dataset: + saved = dataset["data/demo_0/transition"] + for key in snapshot: + torch.testing.assert_close(torch.from_numpy(saved[key][:]), snapshot[key].cpu(), rtol=0, atol=0) + + +def test_terminal_capture_and_return_accounting(task): + env, _ = task + command = env.command_manager.get_term("motion") + capture = env.recorder_manager.get_term("transition") + reward_cfg = env.reward_manager.get_term_cfg("tracking") + term_cfgs = {name: env.termination_manager.get_term_cfg(name) for name in ("failed", "clip_end", "time_out")} + gamma = command.cfg.gamma + try: + env.reward_manager.set_term_cfg( + "tracking", RewardTermCfg(func=lambda env: env.episode_length_buf / env.step_dt, weight=1.0) + ) + env.termination_manager.set_term_cfg( + "failed", + TerminationTermCfg( + func=lambda env: ( + (env.episode_length_buf == 3) & (torch.arange(env.num_envs, device=env.device) % 2 == 0) + ) + ), + ) + env.termination_manager.set_term_cfg( + "time_out", + TerminationTermCfg( + func=lambda env: ( + (env.episode_length_buf == 3) & (torch.arange(env.num_envs, device=env.device) % 2 == 1) + ), + time_out=True, + ), + ) + command.cfg.gamma = 0.9 + env.reset() + command.bin_return.zero_() + command.start_bin.zero_() + for _ in range(3): + env.step(torch.zeros(env.num_envs, 29, device=env.device)) + assert capture.latest["done"].all() + assert capture.latest["term"].sum() == env.num_envs // 2 + assert (capture.latest["ep_step"] == 2).all() + assert (env.episode_length_buf == 0).all() + assert (capture.latest["s1"][:, :3] != command.state()[0]).any() + # Discounted 1,2,3 = 5.23; timeout adds .9**3 * mean(1,2,3)/.1 = 14.58. + torch.testing.assert_close(command.bin_return[0], torch.tensor(0.2 * (5.23 + 19.81) / 2, device=env.device)) + before = command.bin_return.clone() + env.reset() + torch.testing.assert_close(command.bin_return, before, rtol=0, atol=0) + finally: + command.cfg.gamma = gamma + env.reward_manager.set_term_cfg("tracking", reward_cfg) + for name, cfg in term_cfgs.items(): + env.termination_manager.set_term_cfg(name, cfg) diff --git a/tools/environ_docs.py b/tools/environ_docs.py index b1c0141e48d3..858f39d4a2a9 100644 --- a/tools/environ_docs.py +++ b/tools/environ_docs.py @@ -58,7 +58,7 @@ ) # RL libraries listed in a stable order across generated docs. -_RL_LIBRARY_ORDER = ("rl_games", "rsl_rl", "skrl", "sb3", "rlinf", "torchrl") +_RL_LIBRARY_ORDER = ("rl_games", "rsl_rl", "skrl", "sb3", "rlinf", "supermimic", "torchrl") _ALGORITHM_LABELS = {"": "PPO", "ppo": "PPO", "amp": "AMP", "ippo": "IPPO", "mappo": "MAPPO"} # Gym IDs excluded from the training list. The ``-Eval`` suffix marks dedicated @@ -223,7 +223,9 @@ def parse_rl_libraries_from_kwargs(kwargs: dict) -> dict[str, list[str]]: agents: dict[str, set[str]] = collections.defaultdict(set) for value, library, algo_suffix in entries: - algorithm = _ALGORITHM_LABELS.get(algo_suffix.lower()) + algorithm = ( + "MBRL" if library == "supermimic" and not algo_suffix else _ALGORITHM_LABELS.get(algo_suffix.lower()) + ) if algorithm is None: # Registry suffixes can name recipes or preset variants, not algorithms. continue diff --git a/tools/test/test_environ_docs.py b/tools/test/test_environ_docs.py index 48a82681e63d..667b3951b9d6 100644 --- a/tools/test/test_environ_docs.py +++ b/tools/test/test_environ_docs.py @@ -109,11 +109,13 @@ def test_parse_rl_libraries_from_kwargs_handles_multi_agent_and_amp(): "skrl_mappo_cfg_entry_point": "agents:skrl_mappo_cfg.yaml", "skrl_amp_cfg_entry_point": "agents:skrl_amp_cfg.yaml", "rsl_rl_with_symmetry_cfg_entry_point": "agents.rsl_rl_ppo_cfg:RunnerCfg", + "supermimic_cfg_entry_point": "isaaclab_rl.supermimic.supermimic_cfg:SuperMimicCfg", "rl_games_cfg_entry_point_ignored": "nope", } agents = parse_rl_libraries_from_kwargs(kwargs) assert agents["rl_games"] == ["PPO"] assert agents["skrl"] == ["PPO", "AMP", "IPPO", "MAPPO"] + assert agents["supermimic"] == ["MBRL"] def test_parse_rl_libraries_does_not_label_canonical_amp_alias_as_ppo():