From 3e3687d022c0c23947926a154350dde76deb78a8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?B=C3=A1lint=20Hodossy?= Date: Tue, 5 Dec 2023 18:49:32 +0000 Subject: [PATCH 01/92] Add functionality to synchronise rendered terrain to position of the geom --- unity/Runtime/Components/Shapes/MjHeightFieldShape.cs | 5 ++++- unity/Runtime/Tools/MjcfGenerationContext.cs | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs b/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs index 14597256..86b3fee9 100644 --- a/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs +++ b/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs @@ -36,8 +36,11 @@ public class MjHeightFieldShape : IMjShape public void ToMjcf(XmlElement mjcf, Transform transform){ ExportHeightMap(); + if(terrain.transform.parent != transform) Debug.LogWarning($"The terrain of heightfield {transform.name} needs to be parented to the Geom for proper rendering."); + else terrain.transform.localPosition = new Vector3(-HeightMapHeight*HeightMapScale.x / 2, terrain.transform.localPosition.y, -HeightMapWidth*HeightMapScale.z / 2); var scene = MjScene.Instance; var assetName = scene.GenerationContext.AddHeightFieldAsset(this); + mjcf.SetAttribute("hfield", assetName); } @@ -59,7 +62,7 @@ public class MjHeightFieldShape : IMjShape } public Tuple BuildMesh(){ - return Tuple.Create(new Vector3[]{}, new int[]{}); + return null; } public void DebugDraw(Transform transform){ diff --git a/unity/Runtime/Tools/MjcfGenerationContext.cs b/unity/Runtime/Tools/MjcfGenerationContext.cs index 454a6906..86fb9ae7 100644 --- a/unity/Runtime/Tools/MjcfGenerationContext.cs +++ b/unity/Runtime/Tools/MjcfGenerationContext.cs @@ -121,7 +121,7 @@ public class MjcfGenerationContext { mjcf.SetAttribute("content_type", "image/png"); mjcf.SetAttribute("file", hFieldComponent.FullHeightMapPath); - mjcf.SetAttribute("size", MjEngineTool.MakeLocaleInvariant($"{hFieldComponent.HeightMapScale.x*hFieldComponent.HeightMapHeight/2} {hFieldComponent.HeightMapScale.z * hFieldComponent.HeightMapWidth/2} {hFieldComponent.HeightMapScale.y*2} {hFieldComponent.terrain.transform.position.y}")); + mjcf.SetAttribute("size", MjEngineTool.MakeLocaleInvariant($"{hFieldComponent.HeightMapScale.x*hFieldComponent.HeightMapHeight/2} {hFieldComponent.HeightMapScale.z * hFieldComponent.HeightMapWidth/2} {hFieldComponent.HeightMapScale.y} {hFieldComponent.terrain.transform.localPosition.y}")); } } From 2adfd09b875b9a5a2e1bcb7dae70a7b6dba1a8f9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?B=C3=A1lint=20Hodossy?= Date: Thu, 7 Dec 2023 13:34:33 +0000 Subject: [PATCH 02/92] Rebuild scenes on terrain changes for dynamic heighfields --- .../Components/Shapes/MjHeightFieldShape.cs | 46 +++++++++++++++++-- unity/Runtime/Tools/MjcfGenerationContext.cs | 7 +-- 2 files changed, 47 insertions(+), 6 deletions(-) diff --git a/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs b/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs index 86b3fee9..390de7da 100644 --- a/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs +++ b/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs @@ -24,6 +24,7 @@ namespace Mujoco [Serializable] public class MjHeightFieldShape : IMjShape { + [Tooltip("Terrain's heightmap should have a minimum value of zero (fully black).")] public Terrain terrain; [Tooltip("The path, relative to Application.dataPath, where the heightmap data will be save in PNG format.")] @@ -31,16 +32,41 @@ public class MjHeightFieldShape : IMjShape public string FullHeightMapPath => Path.GetFullPath(Path.Combine(Application.dataPath, heightMapExportPath)); public int HeightMapWidth => terrain.terrainData.heightmapTexture.width; - public int HeightMapHeight => terrain.terrainData.heightmapTexture.height; + public int HeightMapLength => terrain.terrainData.heightmapTexture.height; public Vector3 HeightMapScale => terrain.terrainData.heightmapScale; - public void ToMjcf(XmlElement mjcf, Transform transform){ + [Tooltip("At least this many frames will have to pass before the scene is rebuilt with an updated heightmap. Leave as 0 to not update the hfield during simulation. Increasing this can improve performance.")] + public int UpdateLimit; + private int updateCountdown; + + [HideInInspector] + public float MinimumHeight; + [HideInInspector] + public float MaximumHeight; + + public int HeightFieldId { get; private set; } + + public unsafe void ToMjcf(XmlElement mjcf, Transform transform){ ExportHeightMap(); if(terrain.transform.parent != transform) Debug.LogWarning($"The terrain of heightfield {transform.name} needs to be parented to the Geom for proper rendering."); - else terrain.transform.localPosition = new Vector3(-HeightMapHeight*HeightMapScale.x / 2, terrain.transform.localPosition.y, -HeightMapWidth*HeightMapScale.z / 2); + else { + if((terrain.transform.localPosition - new Vector3(-HeightMapLength * HeightMapScale.x / 2, terrain.transform.localPosition.y, -HeightMapWidth * HeightMapScale.z / 2)).magnitude > 0.001) { + Debug.LogWarning($"Terrain of heightfield {transform.name} not aligned with geom. The terrain will be moved to accurately represent the simulated position."); + } + terrain.transform.localPosition = new Vector3(-HeightMapLength*HeightMapScale.x / 2, terrain.transform.localPosition.y, -HeightMapWidth*HeightMapScale.z / 2); + } var scene = MjScene.Instance; var assetName = scene.GenerationContext.AddHeightFieldAsset(this); + scene.postInitEvent += (_,_) => HeightFieldId = MujocoLib.mj_name2id(scene.Model, (int)MujocoLib.mjtObj.mjOBJ_HFIELD, assetName); + + + if(UpdateLimit>0){ + updateCountdown = UpdateLimit; + scene.preUpdateEvent += (_, _) => CountdownUpdateCondition(); + TerrainCallbacks.heightmapChanged += RebuildScene; + } + mjcf.SetAttribute("hfield", assetName); } @@ -52,10 +78,24 @@ public class MjHeightFieldShape : IMjShape RenderTexture.active = terrain.terrainData.heightmapTexture; Texture2D texture = new Texture2D(RenderTexture.active.width, RenderTexture.active.height); texture.ReadPixels(new Rect(0, 0, RenderTexture.active.width, RenderTexture.active.height), 0, 0); + MaximumHeight = texture.GetPixels().Select(c => c.r).Max()*HeightMapScale.y*2; + var minimumHeight = texture.GetPixels().Select(c => c.r).Min()*HeightMapScale.y*2; + if (minimumHeight > 0.0001) Debug.LogWarning("Due to assumptions in MuJoCo heightfields, terrains should have a minimum heightmap value of 0."); RenderTexture.active = null; File.WriteAllBytes(FullHeightMapPath, texture.EncodeToPNG()); } + public void CountdownUpdateCondition() { + if(updateCountdown < 1) return; + updateCountdown -= 1; + } + + public void RebuildScene(Terrain terrain, RectInt heightRegion, bool synched){ + if(updateCountdown > 0) return; + if(!Application.isPlaying || !MjScene.InstanceExists) return; + MjScene.Instance.SceneRecreationAtLateUpdateRequested = true; + updateCountdown = UpdateLimit; + } public Vector4 GetChangeStamp(){ return Vector4.one; diff --git a/unity/Runtime/Tools/MjcfGenerationContext.cs b/unity/Runtime/Tools/MjcfGenerationContext.cs index 86fb9ae7..9a6f3db6 100644 --- a/unity/Runtime/Tools/MjcfGenerationContext.cs +++ b/unity/Runtime/Tools/MjcfGenerationContext.cs @@ -115,13 +115,14 @@ public class MjcfGenerationContext { mjcf.SetAttribute("vertex", vertexPositionsStr.ToString()); } - private static void GenerateHeightFieldMjcf(MjHeightFieldShape hFieldComponent, XmlElement mjcf) { + private static void GenerateHeightFieldMjcf(MjHeightFieldShape hFieldComponent, XmlElement mjcf) { mjcf.SetAttribute("nrow", "0"); mjcf.SetAttribute("ncol", "0"); mjcf.SetAttribute("content_type", "image/png"); mjcf.SetAttribute("file", hFieldComponent.FullHeightMapPath); - - mjcf.SetAttribute("size", MjEngineTool.MakeLocaleInvariant($"{hFieldComponent.HeightMapScale.x*hFieldComponent.HeightMapHeight/2} {hFieldComponent.HeightMapScale.z * hFieldComponent.HeightMapWidth/2} {hFieldComponent.HeightMapScale.y} {hFieldComponent.terrain.transform.localPosition.y}")); + var baseHeight = hFieldComponent.terrain.transform.localPosition.y + hFieldComponent.MinimumHeight; + var heightRange = Mathf.Clamp(hFieldComponent.MaximumHeight - hFieldComponent.MinimumHeight, 0.00001f, Mathf.Infinity); + mjcf.SetAttribute("size", MjEngineTool.MakeLocaleInvariant($"{hFieldComponent.HeightMapScale.x*hFieldComponent.HeightMapLength/2} {hFieldComponent.HeightMapScale.z * hFieldComponent.HeightMapWidth/2} {heightRange} {baseHeight}")); } } From 0b413656b8592239f934feea2c33d50643f63596 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?B=C3=A1lint=20Hodossy?= Date: Thu, 7 Dec 2023 14:02:21 +0000 Subject: [PATCH 03/92] Restrict access to height range properties of the HField shape --- unity/Runtime/Components/Shapes/MjHeightFieldShape.cs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs b/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs index 390de7da..002d082c 100644 --- a/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs +++ b/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs @@ -40,9 +40,9 @@ public class MjHeightFieldShape : IMjShape private int updateCountdown; [HideInInspector] - public float MinimumHeight; + public float MinimumHeight { get; private set; } [HideInInspector] - public float MaximumHeight; + public float MaximumHeight { get; private set; } public int HeightFieldId { get; private set; } From af133f402230f406f888cb415d486f35519d3425 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?B=C3=A1lint=20Hodossy?= Date: Mon, 11 Dec 2023 14:54:55 +0000 Subject: [PATCH 04/92] Fix whitespace placements --- .../Components/Shapes/MjHeightFieldShape.cs | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs b/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs index 002d082c..30226fc9 100644 --- a/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs +++ b/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs @@ -31,7 +31,7 @@ public class MjHeightFieldShape : IMjShape public string heightMapExportPath; public string FullHeightMapPath => Path.GetFullPath(Path.Combine(Application.dataPath, heightMapExportPath)); - public int HeightMapWidth => terrain.terrainData.heightmapTexture.width; + public int HeightMapWidth => terrain.terrainData.heightmapTexture.width; public int HeightMapLength => terrain.terrainData.heightmapTexture.height; public Vector3 HeightMapScale => terrain.terrainData.heightmapScale; @@ -46,7 +46,7 @@ public class MjHeightFieldShape : IMjShape public int HeightFieldId { get; private set; } - public unsafe void ToMjcf(XmlElement mjcf, Transform transform){ + public unsafe void ToMjcf(XmlElement mjcf, Transform transform) { ExportHeightMap(); if(terrain.transform.parent != transform) Debug.LogWarning($"The terrain of heightfield {transform.name} needs to be parented to the Geom for proper rendering."); else { @@ -60,7 +60,7 @@ public class MjHeightFieldShape : IMjShape scene.postInitEvent += (_,_) => HeightFieldId = MujocoLib.mj_name2id(scene.Model, (int)MujocoLib.mjtObj.mjOBJ_HFIELD, assetName); - + if(UpdateLimit>0){ updateCountdown = UpdateLimit; scene.preUpdateEvent += (_, _) => CountdownUpdateCondition(); @@ -70,11 +70,11 @@ public class MjHeightFieldShape : IMjShape mjcf.SetAttribute("hfield", assetName); } - public void FromMjcf(XmlElement mjcf){ + public void FromMjcf(XmlElement mjcf) { } - public void ExportHeightMap(){ + public void ExportHeightMap() { RenderTexture.active = terrain.terrainData.heightmapTexture; Texture2D texture = new Texture2D(RenderTexture.active.width, RenderTexture.active.height); texture.ReadPixels(new Rect(0, 0, RenderTexture.active.width, RenderTexture.active.height), 0, 0); @@ -90,22 +90,22 @@ public class MjHeightFieldShape : IMjShape updateCountdown -= 1; } - public void RebuildScene(Terrain terrain, RectInt heightRegion, bool synched){ + public void RebuildScene(Terrain terrain, RectInt heightRegion, bool synched) { if(updateCountdown > 0) return; if(!Application.isPlaying || !MjScene.InstanceExists) return; MjScene.Instance.SceneRecreationAtLateUpdateRequested = true; updateCountdown = UpdateLimit; } - public Vector4 GetChangeStamp(){ + public Vector4 GetChangeStamp() { return Vector4.one; } - public Tuple BuildMesh(){ + public Tuple BuildMesh() { return null; } - public void DebugDraw(Transform transform){ + public void DebugDraw(Transform transform) { } } } From 315a2f334f8f0da328492a50ac2fc8b0f5b3875c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?B=C3=A1lint=20Hodossy?= Date: Mon, 11 Dec 2023 14:57:40 +0000 Subject: [PATCH 05/92] Fix whitespace in reformatted code from GitHub --- unity/Runtime/Tools/MjcfGenerationContext.cs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/unity/Runtime/Tools/MjcfGenerationContext.cs b/unity/Runtime/Tools/MjcfGenerationContext.cs index 1e75328c..1e3f07ca 100644 --- a/unity/Runtime/Tools/MjcfGenerationContext.cs +++ b/unity/Runtime/Tools/MjcfGenerationContext.cs @@ -125,9 +125,9 @@ public class MjcfGenerationContext { mjcf.SetAttribute( "size", MjEngineTool.MakeLocaleInvariant( - $@"{hFieldComponent.HeightMapScale.x*hFieldComponent.HeightMapLength/2} - {hFieldComponent.HeightMapScale.z * hFieldComponent.HeightMapWidth/2} - {heightRange} + $@"{hFieldComponent.HeightMapScale.x*hFieldComponent.HeightMapLength/2} + {hFieldComponent.HeightMapScale.z * hFieldComponent.HeightMapWidth/2} + {heightRange} {baseHeight}")); } } From d0b505c67b6a22b0dff0da8ede6fc0f5d8c2a2f2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?B=C3=A1lint=20Hodossy?= Date: Tue, 19 Dec 2023 13:54:49 +0000 Subject: [PATCH 06/92] Reformat scripts to better follow guidelines --- .../Components/Shapes/MjHeightFieldShape.cs | 84 +++++++++++-------- unity/Runtime/Tools/MjcfGenerationContext.cs | 25 +++--- 2 files changed, 65 insertions(+), 44 deletions(-) diff --git a/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs b/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs index 30226fc9..46d05fca 100644 --- a/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs +++ b/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs @@ -18,51 +18,66 @@ using System.Linq; using System.Xml; using UnityEngine; -namespace Mujoco -{ +namespace Mujoco { [Serializable] -public class MjHeightFieldShape : IMjShape -{ +public class MjHeightFieldShape : IMjShape { [Tooltip("Terrain's heightmap should have a minimum value of zero (fully black).")] - public Terrain terrain; + public Terrain Terrain; - [Tooltip("The path, relative to Application.dataPath, where the heightmap data will be save in PNG format.")] - public string heightMapExportPath; + [Tooltip("The path, relative to Application.dataPath, where the heightmap " + + "data will be save in PNG format.")] + public string HeightMapExportPath; - public string FullHeightMapPath => Path.GetFullPath(Path.Combine(Application.dataPath, heightMapExportPath)); - public int HeightMapWidth => terrain.terrainData.heightmapTexture.width; - public int HeightMapLength => terrain.terrainData.heightmapTexture.height; - public Vector3 HeightMapScale => terrain.terrainData.heightmapScale; + public string FullHeightMapPath => Path.GetFullPath(Path.Combine(Application.dataPath, + HeightMapExportPath)); - [Tooltip("At least this many frames will have to pass before the scene is rebuilt with an updated heightmap. Leave as 0 to not update the hfield during simulation. Increasing this can improve performance.")] + public int HeightMapWidth => Terrain.terrainData.heightmapTexture.width; + public int HeightMapLength => Terrain.terrainData.heightmapTexture.height; + public Vector3 HeightMapScale => Terrain.terrainData.heightmapScale; + + [Tooltip("At least this many frames will have to pass before the scene is rebuilt with an " + + "updated heightmap. Leave as 0 to not update the hfield during simulation. " + + "Increasing this can improve performance.")] public int UpdateLimit; - private int updateCountdown; - [HideInInspector] + private int _updateCountdown; + + [HideInInspector] public float MinimumHeight { get; private set; } - [HideInInspector] + + [HideInInspector] public float MaximumHeight { get; private set; } public int HeightFieldId { get; private set; } public unsafe void ToMjcf(XmlElement mjcf, Transform transform) { ExportHeightMap(); - if(terrain.transform.parent != transform) Debug.LogWarning($"The terrain of heightfield {transform.name} needs to be parented to the Geom for proper rendering."); + if (Terrain.transform.parent != transform) + Debug.LogWarning( + $"The terrain of heightfield {transform.name} needs to be parented to the Geom " + + "for proper rendering."); else { - if((terrain.transform.localPosition - new Vector3(-HeightMapLength * HeightMapScale.x / 2, terrain.transform.localPosition.y, -HeightMapWidth * HeightMapScale.z / 2)).magnitude > 0.001) { - Debug.LogWarning($"Terrain of heightfield {transform.name} not aligned with geom. The terrain will be moved to accurately represent the simulated position."); + if ((Terrain.transform.localPosition - new Vector3(-HeightMapLength * HeightMapScale.x / 2, + Terrain.transform.localPosition.y, + -HeightMapWidth * HeightMapScale.z / 2)).magnitude > 0.001) { + Debug.LogWarning($"Terrain of heightfield {transform.name} not aligned with geom. The " + + " terrain will be moved to accurately represent the simulated position."); } - terrain.transform.localPosition = new Vector3(-HeightMapLength*HeightMapScale.x / 2, terrain.transform.localPosition.y, -HeightMapWidth*HeightMapScale.z / 2); + Terrain.transform.localPosition = new Vector3(-HeightMapLength * HeightMapScale.x / 2, + Terrain.transform.localPosition.y, + -HeightMapWidth * HeightMapScale.z / 2); } var scene = MjScene.Instance; var assetName = scene.GenerationContext.AddHeightFieldAsset(this); - scene.postInitEvent += (_,_) => HeightFieldId = MujocoLib.mj_name2id(scene.Model, (int)MujocoLib.mjtObj.mjOBJ_HFIELD, assetName); + scene.postInitEvent += (_, _) => + HeightFieldId = + MujocoLib.mj_name2id(scene.Model, (int)MujocoLib.mjtObj.mjOBJ_HFIELD, assetName); - if(UpdateLimit>0){ - updateCountdown = UpdateLimit; + if (UpdateLimit > 0) { + _updateCountdown = UpdateLimit; scene.preUpdateEvent += (_, _) => CountdownUpdateCondition(); TerrainCallbacks.heightmapChanged += RebuildScene; } @@ -71,30 +86,33 @@ public class MjHeightFieldShape : IMjShape } public void FromMjcf(XmlElement mjcf) { - } public void ExportHeightMap() { - RenderTexture.active = terrain.terrainData.heightmapTexture; + RenderTexture.active = Terrain.terrainData.heightmapTexture; Texture2D texture = new Texture2D(RenderTexture.active.width, RenderTexture.active.height); - texture.ReadPixels(new Rect(0, 0, RenderTexture.active.width, RenderTexture.active.height), 0, 0); - MaximumHeight = texture.GetPixels().Select(c => c.r).Max()*HeightMapScale.y*2; - var minimumHeight = texture.GetPixels().Select(c => c.r).Min()*HeightMapScale.y*2; - if (minimumHeight > 0.0001) Debug.LogWarning("Due to assumptions in MuJoCo heightfields, terrains should have a minimum heightmap value of 0."); + texture.ReadPixels(new Rect(0, 0, RenderTexture.active.width, RenderTexture.active.height), + 0, + 0); + MaximumHeight = texture.GetPixels().Select(c => c.r).Max() * HeightMapScale.y * 2; + var minimumHeight = texture.GetPixels().Select(c => c.r).Min() * HeightMapScale.y * 2; + if (minimumHeight > 0.0001) + Debug.LogWarning("Due to assumptions in MuJoCo heightfields, terrains should have a " + + "minimum heightmap value of 0."); RenderTexture.active = null; File.WriteAllBytes(FullHeightMapPath, texture.EncodeToPNG()); } public void CountdownUpdateCondition() { - if(updateCountdown < 1) return; - updateCountdown -= 1; + if (_updateCountdown < 1) return; + _updateCountdown -= 1; } public void RebuildScene(Terrain terrain, RectInt heightRegion, bool synched) { - if(updateCountdown > 0) return; - if(!Application.isPlaying || !MjScene.InstanceExists) return; + if (_updateCountdown > 0) return; + if (!Application.isPlaying || !MjScene.InstanceExists) return; MjScene.Instance.SceneRecreationAtLateUpdateRequested = true; - updateCountdown = UpdateLimit; + _updateCountdown = UpdateLimit; } public Vector4 GetChangeStamp() { diff --git a/unity/Runtime/Tools/MjcfGenerationContext.cs b/unity/Runtime/Tools/MjcfGenerationContext.cs index 1e3f07ca..2c2ba5f2 100644 --- a/unity/Runtime/Tools/MjcfGenerationContext.cs +++ b/unity/Runtime/Tools/MjcfGenerationContext.cs @@ -34,7 +34,9 @@ public class MjcfGenerationContext { private int _nuserSensor; private int _numGeneratedNames = 0; private Dictionary _meshAssets = new Dictionary(); - private Dictionary _hFieldAssets = new Dictionary(); + + private Dictionary _hFieldAssets = + new Dictionary(); public void GenerateMjcf(XmlElement mjcf) { GenerateConfigurationMjcf(mjcf); @@ -98,7 +100,7 @@ public class MjcfGenerationContext { meshMjcf.SetAttribute("name", meshAsset.Value); GenerateMeshMjcf(meshAsset.Key, meshMjcf); } - foreach (var hFieldAsset in _hFieldAssets) { + foreach (var hFieldAsset in _hFieldAssets) { var hFieldMjcf = (XmlElement)assetMjcf.AppendChild(doc.CreateElement("hfield")); hFieldMjcf.SetAttribute("name", hFieldAsset.Value); GenerateHeightFieldMjcf(hFieldAsset.Key, hFieldMjcf); @@ -114,21 +116,22 @@ public class MjcfGenerationContext { } mjcf.SetAttribute("vertex", vertexPositionsStr.ToString()); } - + private static void GenerateHeightFieldMjcf(MjHeightFieldShape hFieldComponent, XmlElement mjcf) { mjcf.SetAttribute("nrow", "0"); mjcf.SetAttribute("ncol", "0"); mjcf.SetAttribute("content_type", "image/png"); mjcf.SetAttribute("file", hFieldComponent.FullHeightMapPath); - var baseHeight = hFieldComponent.terrain.transform.localPosition.y + hFieldComponent.MinimumHeight; - var heightRange = Mathf.Clamp(hFieldComponent.MaximumHeight - hFieldComponent.MinimumHeight, 0.00001f, Mathf.Infinity); + var baseHeight = hFieldComponent.Terrain.transform.localPosition.y + + hFieldComponent.MinimumHeight; + var heightRange = Mathf.Clamp( + hFieldComponent.MaximumHeight - hFieldComponent.MinimumHeight, 0.00001f, Mathf.Infinity); mjcf.SetAttribute( - "size", - MjEngineTool.MakeLocaleInvariant( - $@"{hFieldComponent.HeightMapScale.x*hFieldComponent.HeightMapLength/2} - {hFieldComponent.HeightMapScale.z * hFieldComponent.HeightMapWidth/2} - {heightRange} - {baseHeight}")); + "size", + MjEngineTool.MakeLocaleInvariant( + $@"{hFieldComponent.HeightMapScale.x * hFieldComponent.HeightMapLength / 2} { + hFieldComponent.HeightMapScale.z * hFieldComponent.HeightMapWidth / 2} {heightRange} { + baseHeight}")); } } } From 2c024164eab56ccea4c4785c4c7e27f64254c916 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?B=C3=A1lint=20Hodossy?= Date: Thu, 21 Dec 2023 16:04:25 +0000 Subject: [PATCH 07/92] Edit hfield_data in Unity --- .../Components/Shapes/MjHeightFieldShape.cs | 81 +++++++++++++------ unity/Runtime/Tools/MjcfGenerationContext.cs | 25 ++++-- 2 files changed, 74 insertions(+), 32 deletions(-) diff --git a/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs b/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs index 46d05fca..7d61b457 100644 --- a/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs +++ b/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs @@ -26,9 +26,12 @@ public class MjHeightFieldShape : IMjShape { public Terrain Terrain; [Tooltip("The path, relative to Application.dataPath, where the heightmap " + - "data will be save in PNG format.")] + "data will be saved/exported in PNG format. Leave blank if hfield data should " + + "be set instead programmatically (faster).")] public string HeightMapExportPath; + public bool ExportImage => !string.IsNullOrEmpty(HeightMapExportPath); + public string FullHeightMapPath => Path.GetFullPath(Path.Combine(Application.dataPath, HeightMapExportPath)); @@ -38,57 +41,59 @@ public class MjHeightFieldShape : IMjShape { [Tooltip("At least this many frames will have to pass before the scene is rebuilt with an " + "updated heightmap. Leave as 0 to not update the hfield during simulation. " + - "Increasing this can improve performance.")] + "If nonzero, increasing this can improve performance.")] public int UpdateLimit; private int _updateCountdown; - [HideInInspector] + [HideInInspector] public float MinimumHeight { get; private set; } - [HideInInspector] + [HideInInspector] public float MaximumHeight { get; private set; } public int HeightFieldId { get; private set; } public unsafe void ToMjcf(XmlElement mjcf, Transform transform) { - ExportHeightMap(); if (Terrain.transform.parent != transform) Debug.LogWarning( $"The terrain of heightfield {transform.name} needs to be parented to the Geom " + "for proper rendering."); else { - if ((Terrain.transform.localPosition - new Vector3(-HeightMapLength * HeightMapScale.x / 2, + if ((Terrain.transform.localPosition - new Vector3( + -(HeightMapLength - 1) * HeightMapScale.x / 2f, Terrain.transform.localPosition.y, - -HeightMapWidth * HeightMapScale.z / 2)).magnitude > 0.001) { + -(HeightMapWidth - 1) * HeightMapScale.z / 2)).magnitude > 0.001) { Debug.LogWarning($"Terrain of heightfield {transform.name} not aligned with geom. The " + " terrain will be moved to accurately represent the simulated position."); } - Terrain.transform.localPosition = new Vector3(-HeightMapLength * HeightMapScale.x / 2, + Terrain.transform.localPosition = new Vector3(-(HeightMapLength - 1) * HeightMapScale.x / 2, Terrain.transform.localPosition.y, - -HeightMapWidth * HeightMapScale.z / 2); + -(HeightMapWidth - 1) * HeightMapScale.z / 2); } var scene = MjScene.Instance; var assetName = scene.GenerationContext.AddHeightFieldAsset(this); - scene.postInitEvent += (_, _) => - HeightFieldId = - MujocoLib.mj_name2id(scene.Model, (int)MujocoLib.mjtObj.mjOBJ_HFIELD, assetName); - + if (Application.isPlaying) { + scene.postInitEvent += (_, _) => + HeightFieldId = + MujocoLib.mj_name2id(scene.Model, (int)MujocoLib.mjtObj.mjOBJ_HFIELD, assetName); + } if (UpdateLimit > 0) { _updateCountdown = UpdateLimit; - scene.preUpdateEvent += (_, _) => CountdownUpdateCondition(); - TerrainCallbacks.heightmapChanged += RebuildScene; + if (UpdateLimit > 1) scene.preUpdateEvent += (_, _) => CountdownUpdateCondition(); + TerrainCallbacks.heightmapChanged += (_, _, _) => RebuildHeightField(); } mjcf.SetAttribute("hfield", assetName); + PrepareHeightMap(); } public void FromMjcf(XmlElement mjcf) { } - public void ExportHeightMap() { + public void PrepareHeightMap() { RenderTexture.active = Terrain.terrainData.heightmapTexture; Texture2D texture = new Texture2D(RenderTexture.active.width, RenderTexture.active.height); texture.ReadPixels(new Rect(0, 0, RenderTexture.active.width, RenderTexture.active.height), @@ -96,11 +101,32 @@ public class MjHeightFieldShape : IMjShape { 0); MaximumHeight = texture.GetPixels().Select(c => c.r).Max() * HeightMapScale.y * 2; var minimumHeight = texture.GetPixels().Select(c => c.r).Min() * HeightMapScale.y * 2; - if (minimumHeight > 0.0001) - Debug.LogWarning("Due to assumptions in MuJoCo heightfields, terrains should have a " + - "minimum heightmap value of 0."); + RenderTexture.active = null; - File.WriteAllBytes(FullHeightMapPath, texture.EncodeToPNG()); + if (ExportImage) { + if (minimumHeight > 0.0001) + Debug.LogWarning("Due to assumptions in MuJoCo heightfields, terrains should have a " + + "minimum heightmap value of 0."); + File.WriteAllBytes(FullHeightMapPath, texture.EncodeToPNG()); + } else if (Application.isPlaying) { + MjScene.Instance.postInitEvent += (_, _) => UpdateHeightFieldData(); + } + } + + public unsafe void UpdateHeightFieldData() { + RenderTexture.active = Terrain.terrainData.heightmapTexture; + Texture2D texture = new Texture2D(RenderTexture.active.width, RenderTexture.active.height); + texture.ReadPixels(new Rect(0, 0, RenderTexture.active.width, RenderTexture.active.height), + 0, + 0); + RenderTexture.active = null; + + float[] curData = texture.GetPixels(0, 0, texture.width, texture.height) + .Select(c => c.r * 2).ToArray(); + int adr = MjScene.Instance.Model->hfield_adr[HeightFieldId]; + for (int i = 0; i < curData.Length; i++) { + MjScene.Instance.Model->hfield_data[adr + i] = curData[i]; + } } public void CountdownUpdateCondition() { @@ -108,11 +134,18 @@ public class MjHeightFieldShape : IMjShape { _updateCountdown -= 1; } - public void RebuildScene(Terrain terrain, RectInt heightRegion, bool synched) { - if (_updateCountdown > 0) return; + public void RebuildHeightField() { + // The update rate limiting countdown goes from UpdateLimit + 1 to 1, since if the Update + // limit is at 1, we can update on every frame and don't need a countdown. + if (_updateCountdown > 1) return; if (!Application.isPlaying || !MjScene.InstanceExists) return; - MjScene.Instance.SceneRecreationAtLateUpdateRequested = true; - _updateCountdown = UpdateLimit; + if (ExportImage) { + // If we export an image, it needs to be read by the compiler so we might as well rebuild the scene. + MjScene.Instance.SceneRecreationAtLateUpdateRequested = true; + } else { + UpdateHeightFieldData(); + } + _updateCountdown = UpdateLimit + 1; } public Vector4 GetChangeStamp() { diff --git a/unity/Runtime/Tools/MjcfGenerationContext.cs b/unity/Runtime/Tools/MjcfGenerationContext.cs index 2c2ba5f2..09b2abcc 100644 --- a/unity/Runtime/Tools/MjcfGenerationContext.cs +++ b/unity/Runtime/Tools/MjcfGenerationContext.cs @@ -118,20 +118,29 @@ public class MjcfGenerationContext { } private static void GenerateHeightFieldMjcf(MjHeightFieldShape hFieldComponent, XmlElement mjcf) { - mjcf.SetAttribute("nrow", "0"); - mjcf.SetAttribute("ncol", "0"); - mjcf.SetAttribute("content_type", "image/png"); - mjcf.SetAttribute("file", hFieldComponent.FullHeightMapPath); + if (hFieldComponent.ExportImage) { + mjcf.SetAttribute("content_type", "image/png"); + mjcf.SetAttribute("file", hFieldComponent.FullHeightMapPath); + mjcf.SetAttribute("nrow", "0"); + mjcf.SetAttribute("ncol", "0"); + } else { + mjcf.SetAttribute("nrow", hFieldComponent.HeightMapLength.ToString()); + mjcf.SetAttribute("ncol", hFieldComponent.HeightMapWidth.ToString()); + } + var baseHeight = hFieldComponent.Terrain.transform.localPosition.y + hFieldComponent.MinimumHeight; var heightRange = Mathf.Clamp( - hFieldComponent.MaximumHeight - hFieldComponent.MinimumHeight, 0.00001f, Mathf.Infinity); + hFieldComponent.MaximumHeight - hFieldComponent.MinimumHeight, + 0.00001f, + Mathf.Infinity); mjcf.SetAttribute( "size", MjEngineTool.MakeLocaleInvariant( - $@"{hFieldComponent.HeightMapScale.x * hFieldComponent.HeightMapLength / 2} { - hFieldComponent.HeightMapScale.z * hFieldComponent.HeightMapWidth / 2} {heightRange} { - baseHeight}")); + $@"{hFieldComponent.HeightMapScale.x * (hFieldComponent.HeightMapLength - 1) / 2f} { + hFieldComponent.HeightMapScale.z * (hFieldComponent.HeightMapWidth - 1) / 2f} { + heightRange} { + baseHeight}")); } } } From 3c1907e678f7c07799870ba2b148518e7433f772 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Tue, 16 Jan 2024 08:53:33 -0800 Subject: [PATCH 08/92] Fix rendering of links in documentation. PiperOrigin-RevId: 598855374 Change-Id: I19c1ef73d23f0c2fedb8f0a6754bb9edd91f4c20 --- doc/XMLreference.rst | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/doc/XMLreference.rst b/doc/XMLreference.rst index 6100d9c1..072aaab5 100644 --- a/doc/XMLreference.rst +++ b/doc/XMLreference.rst @@ -311,15 +311,16 @@ any effect. The settings here are global and apply to the entire model. - All materials are discarded. - All textures are discarded. - - All geoms with :ref:`contype`=:ref:`conaffinity`=0 are discarded, if they - are not referenced in another MJCF element. If a discarded geom was used for inferring body inertia, an explicit - :ref:`inertial` element is added to the body. + - All geoms with :ref:`contype` |-| = |-| :ref:`conaffinity` |-| =0 are + discarded, if they are not referenced in another MJCF element. If a discarded geom was used for inferring body + inertia, an explicit :ref:`inertial` element is added to the body. - All meshes which are not referenced by any geom (in particular those discarded above) are discarded. - The resulting compiled model will have exactly the same dynamics as the original model, with the exception of - raycasting, as used for example by :ref:`rangefinder`, since raycasting reports distances to - visual geoms. When visualizing models compiled with this flag, it is important to remember that colliding geoms are - often placed in a :ref:`group` which is invisible by default. + The resulting compiled model will have exactly the same dynamics as the original model. The only engine-level + computation which might change is the output of :ref:`raycasting` computations, as used for example by + :ref:`rangefinder` sensors, since raycasting reports distances to visual geoms. When visualizing + models compiled with this flag, it is important to remember that collision geoms are often placed in a + :ref:`group` which is invisible by default. .. _compiler-convexhull: From 80674149ce486f039288e4b72d379680d1f48716 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Tue, 16 Jan 2024 09:31:15 -0800 Subject: [PATCH 09/92] Allow pinning flexes to non-simple bodies. Fixes #1270. PiperOrigin-RevId: 598865882 Change-Id: I105ff1910a7d67f5b7bcb46bcef0e5e388ec08ea --- plugin/elasticity/elasticity.h | 39 +++++++++++++++++++++++----------- plugin/elasticity/membrane.cc | 12 +++++++++-- plugin/elasticity/membrane.h | 1 + plugin/elasticity/solid.cc | 12 +++++++++-- plugin/elasticity/solid.h | 1 + 5 files changed, 49 insertions(+), 16 deletions(-) diff --git a/plugin/elasticity/elasticity.h b/plugin/elasticity/elasticity.h index 00074a6e..fa70e356 100644 --- a/plugin/elasticity/elasticity.h +++ b/plugin/elasticity/elasticity.h @@ -86,13 +86,14 @@ void inline GradSquaredLengths(mjtNum gradient[T::kNumEdges][2][3], } template -inline void ComputeForce(mjtNum* qfrc_passive, +inline void ComputeForce(std::vector& qfrc_passive, const std::vector& elements, const std::vector& metric, const std::vector& elongationglob, const mjModel* m, - const int* vertbodyid, const mjtNum* xpos) { + mju_zero(qfrc_passive.data(), qfrc_passive.size()); + for (int t = 0; t < elements.size(); t++) { const int* v = elements[t].vertices; @@ -118,7 +119,7 @@ inline void ComputeForce(mjtNum* qfrc_passive, for (int ed2 = 0; ed2 < T::kNumEdges; ed2++) { for (int i = 0; i < 2; i++) { for (int x = 0; x < 3; x++) { - force[3 * T::edge[ed2][i] + x] += + force[3 * T::edge[ed2][i] + x] -= elongation[ed1] * gradient[ed2][i][x] * metric[offset * t + T::kNumEdges * ed1 + ed2]; } @@ -128,17 +129,31 @@ inline void ComputeForce(mjtNum* qfrc_passive, // insert into global force for (int i = 0; i < T::kNumVerts; i++) { - int body_dofnum = 3; - int body_dofadr = 3*v[i]; - if (vertbodyid) { - body_dofnum = m->body_dofnum[vertbodyid[v[i]]]; - body_dofadr = m->body_dofadr[vertbodyid[v[i]]]; - if (body_dofnum && m->body_simple[vertbodyid[v[i]]] != 2) { - mju_error("Non-simple or non-static bodies are not yet supported"); - } + for (int x = 0; x < 3; x++) { + qfrc_passive[3*v[i]+x] += force[3*i+x]; } + } + } +} + +// add flex force to degrees of freedom +inline void AddFlexForce(mjtNum* qfrc, + const std::vector& force, + const mjModel* m, mjData* d, + const mjtNum* xpos, + int f0) { + int* bodyid = m->flex_vertbodyid + m->flex_vertadr[f0]; + + for (int v = 0; v < m->flex_vertnum[f0]; v++) { + int bid = bodyid[v]; + if (m->body_simple[bid] != 2) { + // this should only occur for pinned flex vertices + mj_applyFT(m, d, force.data() + 3*v, 0, xpos + 3*v, bid, qfrc); + } else { + int body_dofnum = m->body_dofnum[bid]; + int body_dofadr = m->body_dofadr[bid]; for (int x = 0; x < body_dofnum; x++) { - qfrc_passive[body_dofadr+x] -= force[3*i+x]; + qfrc[body_dofadr+x] += force[3*v+x]; } } } diff --git a/plugin/elasticity/membrane.cc b/plugin/elasticity/membrane.cc index 51283099..f2f88009 100644 --- a/plugin/elasticity/membrane.cc +++ b/plugin/elasticity/membrane.cc @@ -117,6 +117,7 @@ Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu, for (int j = 0; j < m->flex_vertnum[i]; j++) { if (m->flex_vertbodyid[m->flex_vertadr[i]+j] == i0) { f0 = i; + nv = m->flex_vertnum[f0]; } } } @@ -168,6 +169,7 @@ Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu, deformed.assign(ne, 0); previous.assign(ne, 0); elongation.assign(ne, 0); + force.assign(3*nv, 0); // compute edge lengths at equilibrium (m->flexedge_length0 not yet available) UpdateSquaredLengths(reference, edges, body_pos); @@ -198,11 +200,17 @@ void Membrane::Compute(const mjModel* m, mjData* d, int instance) { // compute gradient of elastic energy and insert into passive force int flex_vertadr = f0 < 0 ? -1 : m->flex_vertadr[f0]; - int* bodyid = f0 < 0 ? nullptr : m->flex_vertbodyid + flex_vertadr; mjtNum* xpos = f0 < 0 ? d->xpos + 3*i0 : d->flexvert_xpos + 3*flex_vertadr; mjtNum* qfrc = d->qfrc_passive + (f0 < 0 ? m->body_dofadr[i0] : 0); - ComputeForce(qfrc, elements, metric, elongation, m, bodyid, xpos); + ComputeForce(force, elements, metric, elongation, m, xpos); + + // insert into passive force + if (f0 < 0) { + mju_addTo(qfrc, force.data(), force.size()); + } else { + AddFlexForce(qfrc, force, m, d, xpos, f0); + } // update stored lengths if (kD > 0) { diff --git a/plugin/elasticity/membrane.h b/plugin/elasticity/membrane.h index 4b39d426..2b9a23da 100644 --- a/plugin/elasticity/membrane.h +++ b/plugin/elasticity/membrane.h @@ -57,6 +57,7 @@ class Membrane { std::vector deformed; // deformed lengths (ne x 1) std::vector previous; // previous-step lengths (ne x 1) std::vector elongation; // edge elongation (ne x 1) + std::vector force; // force at all vertices (nv x 3) mjtNum damping; mjtNum thickness; diff --git a/plugin/elasticity/solid.cc b/plugin/elasticity/solid.cc index 35516a89..5660e144 100644 --- a/plugin/elasticity/solid.cc +++ b/plugin/elasticity/solid.cc @@ -122,6 +122,7 @@ Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E, for (int j = 0; j < m->flex_vertnum[i]; j++) { if (m->flex_vertbodyid[m->flex_vertadr[i]+j] == i0) { f0 = i; + nv = m->flex_vertnum[f0]; } } } @@ -172,6 +173,7 @@ Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E, deformed.assign(ne, 0); previous.assign(ne, 0); elongation.assign(ne, 0); + force.assign(3*nv, 0); // compute edge lengths at equilibrium (m->flexedge_length0 not yet available) UpdateSquaredLengths(reference, edges, body_pos); @@ -202,11 +204,17 @@ void Solid::Compute(const mjModel* m, mjData* d, int instance) { // compute gradient of elastic energy and insert into passive force int flex_vertadr = f0 < 0 ? -1 : m->flex_vertadr[f0]; - int* bodyid = f0 < 0 ? nullptr : m->flex_vertbodyid + flex_vertadr; mjtNum* xpos = f0 < 0 ? d->xpos + 3*i0 : d->flexvert_xpos + 3*flex_vertadr; mjtNum* qfrc = d->qfrc_passive + (f0 < 0 ? m->body_dofadr[i0] : 0); - ComputeForce(qfrc, elements, metric, elongation, m, bodyid, xpos); + ComputeForce(force, elements, metric, elongation, m, xpos); + + // insert into passive force + if (f0 < 0) { + mju_addTo(qfrc, force.data(), force.size()); + } else { + AddFlexForce(qfrc, force, m, d, xpos, f0); + } // update stored lengths if (kD > 0) { diff --git a/plugin/elasticity/solid.h b/plugin/elasticity/solid.h index ae1f2fe5..acbd99a8 100644 --- a/plugin/elasticity/solid.h +++ b/plugin/elasticity/solid.h @@ -55,6 +55,7 @@ class Solid { std::vector deformed; // deformed lengths (ne x 1) std::vector previous; // previous-step lengths (ne x 1) std::vector elongation; // edge elongation (ne x 1) + std::vector force; // force at all vertices (nv x 3) mjtNum damping; From 45208a5fe06bd8b80179dcaba2c68d61cff0c331 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Tue, 16 Jan 2024 10:03:28 -0800 Subject: [PATCH 10/92] Update Sphinx from 4.5.0 to 5.3.0 fixing broken docs build. PiperOrigin-RevId: 598875231 Change-Id: I3295e9aa94087d44a695c7d380bbf9691a3d82b3 --- doc/requirements.txt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/doc/requirements.txt b/doc/requirements.txt index ae4efc06..b382897d 100644 --- a/doc/requirements.txt +++ b/doc/requirements.txt @@ -1,4 +1,4 @@ -Sphinx==4.5.0 +Sphinx==5.3.0 furo==2022.9.29 sphinxcontrib-bibtex==2.6.1 sphinxcontrib-katex==0.9.4 @@ -11,7 +11,7 @@ nbsphinx==0.9.1 pandoc==1.1.0 pygments==2.15.0 jq==1.4.1 -Jinja2==2.11.3 +Jinja2~=3.0 wheel # see https://github.com/aws/aws-sam-cli/issues/3661 regarding markupsafe markupsafe==2.0.1 From ee922be393d790fa56d18579a87bd7bb7bf8dfef Mon Sep 17 00:00:00 2001 From: Erik Frey Date: Tue, 16 Jan 2024 16:17:06 -0800 Subject: [PATCH 11/92] Benchmark improvements. - Use multiple devices if present - FLAGS more agnostic to diverse platforms PiperOrigin-RevId: 598990299 Change-Id: Ifbd454ca521a788324f2b62f1c79e7d4b8305a30 --- mjx/mujoco/mjx/benchmark/benchmark.py | 91 ++++++++++----------------- 1 file changed, 34 insertions(+), 57 deletions(-) diff --git a/mjx/mujoco/mjx/benchmark/benchmark.py b/mjx/mujoco/mjx/benchmark/benchmark.py index c3145f57..ca29a7e1 100644 --- a/mjx/mujoco/mjx/benchmark/benchmark.py +++ b/mjx/mujoco/mjx/benchmark/benchmark.py @@ -27,76 +27,51 @@ from mujoco import mjx FLAGS = flags.FLAGS -_PATHS = { - 'humanoid': 'benchmark/model/humanoid/humanoid.xml', - 'barkour': 'benchmark/model/barkour_v0/assets/barkour_v0_mjx.xml', - 'shadow_hand': 'benchmark/model/shadow_hand/scene_right.xml', -} - -_BATCH_SIZE = { - ('barkour', 'tpu_v5e'): 1024, - ('humanoid', 'tpu_v5e'): 1024, - ('shadow_hand', 'tpu_v5e'): 1024, - ('barkour', 'gpu_a100'): 8192, - ('humanoid', 'gpu_a100'): 8192, - ('shadow_hand', 'gpu_a100'): 4096, - ('barkour', 'cpu'): 64, - ('humanoid', 'cpu'): 64, - ('shadow_hand', 'cpu'): 64, -} - -_SOLVER_CONFIG = { - ('barkour', 'tpu_v5e'): (mujoco.mjtSolver.mjSOL_CG, 4, 6), - ('humanoid', 'tpu_v5e'): (mujoco.mjtSolver.mjSOL_CG, 6, 6), - ('shadow_hand', 'tpu_v5e'): (mujoco.mjtSolver.mjSOL_CG, 8, 6), - ('humanoid', 'gpu_a100'): (mujoco.mjtSolver.mjSOL_NEWTON, 1, 4), - ('barkour', 'gpu_a100'): (mujoco.mjtSolver.mjSOL_NEWTON, 1, 4), - ('shadow_hand', 'gpu_a100'): (mujoco.mjtSolver.mjSOL_NEWTON, 1, 4), - ('barkour', 'cpu'): (mujoco.mjtSolver.mjSOL_NEWTON, 1, 4), - ('humanoid', 'cpu'): (mujoco.mjtSolver.mjSOL_NEWTON, 1, 4), - ('shadow_hand', 'cpu'): (mujoco.mjtSolver.mjSOL_NEWTON, 1, 4), -} +flags.DEFINE_string('mjcf', None, 'path to model', required=True) +flags.DEFINE_integer('step_count', 1000, 'number of steps per rollout') +flags.DEFINE_integer('batch_size', 1024, 'number of parallel rollouts') +flags.DEFINE_integer('unroll', 1, 'loop unroll length') +flags.DEFINE_enum('solver', 'cg', ['cg', 'newton'], 'constraint solver') +flags.DEFINE_integer('iterations', 1, 'number of solver iterations') +flags.DEFINE_integer('ls_iterations', 4, 'number of linesearch iterations') -flags.DEFINE_string('model', 'humanoid', 'Model to benchmark') -flags.DEFINE_enum('device', 'cpu', ('cpu', 'tpu_v5e', 'gpu_a100'), - 'Device benchmark is running on') - - -def _measure_fn(state, init_fn, step_fn, batch_size: int = 1024) -> float: +def _measure(state, init_fn, step_fn) -> float: """Reports jit time and op time for a function.""" - step_count = 100 if FLAGS.device == 'cpu' else 1000 - - @jax.jit + @jax.pmap def run_batch(seed: jp.ndarray): + batch_size = FLAGS.batch_size // jax.device_count() rngs = jax.random.split(jax.random.PRNGKey(seed), batch_size) - init_state = jax.vmap(init_fn)(rngs) + state = jax.vmap(init_fn)(rngs) @jax.vmap - def run(state): - def step(state, _): - state = step_fn(state) - return state, () + def step(state, _): + state = step_fn(state) + return state, None - return jax.lax.scan(step, state, (), length=step_count) - - return run(init_state) + state, _ = jax.lax.scan( + step, state, None, length=FLAGS.step_count, unroll=FLAGS.unroll + ) + return state # run once to jit beg = time.perf_counter() - jax.tree_util.tree_map(lambda x: x.block_until_ready(), run_batch(0)) + seed = 0 + seeds = jp.arange(seed, seed + jax.device_count(), dtype=int) + jax.tree_util.tree_map(lambda x: x.block_until_ready(), run_batch(seeds)) first_t = time.perf_counter() - beg times = [] while state: + seed += jax.device_count() + seeds = jp.arange(seed, seed + jax.device_count(), dtype=int) beg = time.perf_counter() - batch = run_batch(jp.array(len(times))) - jax.tree_util.tree_map(lambda x: x.block_until_ready(), batch) + jax.tree_util.tree_map(lambda x: x.block_until_ready(), run_batch(seeds)) times.append(time.perf_counter() - beg) op_time = jp.mean(jp.array(times)) - batch_sps = batch_size * step_count / op_time + batch_sps = FLAGS.batch_size * FLAGS.step_count / op_time state.counters['jit_time'] = first_t - op_time state.counters['batch_sps'] = batch_sps @@ -106,11 +81,14 @@ def _measure_fn(state, init_fn, step_fn, batch_size: int = 1024) -> float: def _run(state: benchmark.State): """Benchmark a model.""" - f = epath.resource_path('mujoco.mjx') / _PATHS[FLAGS.model] + f = epath.resource_path('mujoco.mjx') / 'benchmark/model' / FLAGS.mjcf m = mujoco.MjModel.from_xml_path(f.as_posix()) - m.opt.solver, m.opt.iterations, m.opt.ls_iterations = _SOLVER_CONFIG[ - (FLAGS.model, FLAGS.device) - ] + m.opt.solver = { + 'cg': mujoco.mjtSolver.mjSOL_CG, + 'newton': mujoco.mjtSolver.mjSOL_NEWTON, + }[FLAGS.solver.lower()] + m.opt.iterations = FLAGS.iterations + m.opt.ls_iterations = FLAGS.ls_iterations m = mjx.device_put(m) def init(rng): @@ -122,11 +100,10 @@ def _run(state: benchmark.State): def step(d): return mjx.step(m, d) - batch_size = _BATCH_SIZE[(FLAGS.model, FLAGS.device)] - _measure_fn(state, init, step, batch_size=batch_size) + _measure(state, init, step) if __name__ == '__main__': FLAGS(sys.argv) - benchmark.register(_run, name=FLAGS.model + '_' + FLAGS.device) + benchmark.register(_run, name=sys.argv[0].split('/')[-1]) benchmark.main() From a02fc405af8da2589b22cbc1ac1b4f5ea783e5ba Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Tue, 16 Jan 2024 16:28:19 -0800 Subject: [PATCH 12/92] Add naive ray mesh implementation. PiperOrigin-RevId: 598993152 Change-Id: I708505bf20a89d5a4840451a11b99771195003e6 --- doc/changelog.rst | 2 +- mjx/mujoco/mjx/_src/ray.py | 87 +++++++++++++++++++++++++++++--- mjx/mujoco/mjx/_src/ray_test.py | 37 ++++++++++++++ mjx/mujoco/mjx/_src/types.py | 14 +++++ mjx/mujoco/mjx/test_data/ray.xml | 2 + 5 files changed, 133 insertions(+), 9 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index fb886273..eb9aeab0 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -15,7 +15,7 @@ MJX 2. Added :ref:`dyntype` ``filterexact``. 3. Added :at:`site` transmission. 4. Updated MJX colab tutorial with more stable quadruped environment. -5. Added ``mjx.ray`` which mirrors :ref:`mj_ray` for planes, spheres, capsules, and boxes. +5. Added ``mjx.ray`` which mirrors :ref:`mj_ray` for planes, spheres, capsules, boxes, and meshes. Bug fixes ^^^^^^^^^ diff --git a/mjx/mujoco/mjx/_src/ray.py b/mjx/mujoco/mjx/_src/ray.py index d093e926..a9413b21 100644 --- a/mjx/mujoco/mjx/_src/ray.py +++ b/mjx/mujoco/mjx/_src/ray.py @@ -19,6 +19,7 @@ from typing import Sequence, Tuple import jax from jax import numpy as jp import mujoco +from mujoco.mjx._src import math # pylint: disable=g-importing-member from mujoco.mjx._src.types import Data from mujoco.mjx._src.types import GeomType @@ -129,14 +130,79 @@ def _ray_box( return jp.min(jp.where(valid, x, jp.inf)) -def _ray_mesh( - size: jax.Array, +def _ray_triangle( + vert: jax.Array, pnt: jax.Array, vec: jax.Array, + b0: jax.Array, + b1: jax.Array, ) -> jax.Array: - """Returns the distance at which a ray intersects with a mesh.""" - del size, pnt, vec - raise NotImplementedError("ray <> mesh not implemented yet") + """Returns the distance at which a ray intersects with a triangle.""" + # project difference vectors in ray normal plane + planar = jp.dot(jp.array([b0, b1]), (vert - pnt).T) + + # determine if origin is inside planar projection of triangle + # A = (p0-p2, p1-p2), b = -p2, solve A*t = b + A = jp.array( # pylint: disable=invalid-name + [planar[:, 0] - planar[:, 2], planar[:, 1] - planar[:, 2]] + ).T.flatten() + b = -planar[:, 2] + det = A[0] * A[3] - A[1] * A[2] + valid = jp.abs(det) >= mujoco.mjMINVAL + + t0 = (A[3] * b[0] - A[1] * b[1]) / det + t1 = (-A[2] * b[0] + A[0] * b[1]) / det + valid &= (t0 >= 0) & (t1 >= 0) & (t0 + t1 <= 1) + + # intersect ray with plane of triangle + nrm = jp.cross(vert[0] - vert[2], vert[1] - vert[2]) + denom = jp.dot(vec, nrm) + valid &= jp.abs(denom) >= mujoco.mjMINVAL + + dist = jp.where(valid, -jp.dot(pnt - vert[2], nrm) / denom, jp.inf) + + return dist + + +def _ray_mesh( + m: Model, + geom_id: np.ndarray, + unused_size: jax.Array, + pnt: jax.Array, + vec: jax.Array, +) -> Tuple[jax.Array, jax.Array]: + """Returns the best distance and geom_id for ray mesh intersections.""" + data_id = m.geom_dataid[geom_id] + + ray_basis = lambda x: math.orthogonals(math.normalize(x)) + b0, b1 = jax.vmap(ray_basis)(vec) + + faceadr = np.append(m.mesh_faceadr, m.nmeshface) + vertadr = np.append(m.mesh_vertadr, m.nmeshvert) + + dists = [] + for i, id_ in enumerate(data_id): + face = m.mesh_face[faceadr[id_] : faceadr[id_ + 1]] + vert = m.mesh_vert[vertadr[id_] : vertadr[id_ + 1]] + dist = jax.vmap(_ray_triangle, in_axes=(0, None, None, None, None))( + vert[face], pnt[i], vec[i], b0[i], b1[i] + ) + dists.append(dist) + + # map the triangle id to data id + tri_id = np.append(0, (faceadr[data_id + 1] - faceadr[data_id]).cumsum()) + tri_data_id = np.zeros(tri_id[-1], dtype=np.int32) + tri_data_id[tri_id[:-1]] = 1 + tri_data_id = tri_data_id.cumsum() - 1 + + dists = jp.concatenate(dists) + min_id = jp.argmin(dists) + # Grab the best distance amongst all meshes, bypassing the argmin in `ray`. + # This avoids having to compute the best distance per mesh. + dist = dists[min_id, None] + id_ = jp.array(geom_id)[jp.array(tri_data_id)[min_id], None] + + return dist, id_ _RAY_FUNC = { @@ -144,7 +210,7 @@ _RAY_FUNC = { GeomType.SPHERE: _ray_sphere, GeomType.CAPSULE: _ray_capsule, GeomType.BOX: _ray_box, - # GeomType.MESH: _ray_mesh, + GeomType.MESH: _ray_mesh, } @@ -192,8 +258,13 @@ def ray( if id_.size == 0: continue - size, pnt, vec = m.geom_size[id_], geom_pnts[id_], geom_vecs[id_] - dist = jax.vmap(fn)(size, pnt, vec) + args = m.geom_size[id_], geom_pnts[id_], geom_vecs[id_] + + if geom_type == GeomType.MESH: + dist, id_ = fn(m, id_, *args) + else: + dist = jax.vmap(fn)(*args) + dists, ids = dists + [dist], ids + [id_] if not ids: diff --git a/mjx/mujoco/mjx/_src/ray_test.py b/mjx/mujoco/mjx/_src/ray_test.py index 0556c9a3..f89f0d87 100644 --- a/mjx/mujoco/mjx/_src/ray_test.py +++ b/mjx/mujoco/mjx/_src/ray_test.py @@ -144,6 +144,43 @@ class RayTest(absltest.TestCase): mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused) _assert_eq(dist, mj_dist, 'dist') + def test_ray_mesh(self): + """Tests MJX ray<>mesh matches MuJoCo.""" + m = test_util.load_test_file('ray.xml') + d = mujoco.MjData(m) + mujoco.mj_forward(m, d) + mx, dx = mjx.put_model(m), mjx.put_data(m, d) + + # look at the tetrahedron + pnt, vec = jp.array([2.0, 2.0, 2.0]), -jp.array([ + 1.0, + 1.0, + 1.0, + ]) + vec /= jp.linalg.norm(vec) + dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec) + _assert_eq(geomid, 4, 'geom_id') + + pnt, vec, geomid = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32) + mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, geomid) + _assert_eq(geomid, 4, 'geom_id') + _assert_eq(dist, mj_dist, 'dist-tetrahedron') + + # look at the dodecahedron + pnt, vec = jp.array([4.0, 2.0, 2.0]), -jp.array([ + 2.0, + 1.0, + 1.0, + ]) + vec /= jp.linalg.norm(vec) + dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec) + _assert_eq(geomid, 5, 'geom_id') + + pnt, vec, geomid = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32) + mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, geomid) + _assert_eq(geomid, 5, 'geom_id') + _assert_eq(dist, mj_dist, 'dist-dodecahedron') + def test_ray_geomgroup(self): """Tests ray geomgroup filter.""" m = test_util.load_test_file('ray.xml') diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 79cecaf5..1527b583 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -264,6 +264,8 @@ class Model(PyTreeNode): ngeom: number of geoms nsite: number of sites nmesh: number of meshes + nmeshvert: number of vertices in all meshes + nmeshface: number of triangular faces in all meshes nmat: number of materials npair: number of predefined geom pairs nexclude: number of excluded geom pairs @@ -321,6 +323,7 @@ class Model(PyTreeNode): geom_conaffinity: geom contact affinity (ngeom,) geom_condim: contact dimensionality (1, 3, 4, 6) (ngeom,) geom_bodyid: id of geom's body (ngeom,) + geom_dataid: id of geom's mesh/hfield; -1: none (ngeom,) geom_group: group for visibility (ngeom,) geom_matid: material id for rendering (ngeom,) geom_priority: geom contact priority (ngeom,) @@ -338,6 +341,10 @@ class Model(PyTreeNode): site_pos: local position offset rel. to body (nsite, 3) site_quat: local orientation offset rel. to body (nsite, 4) mat_rgba: rgba (nmat, 4) + mesh_vertadr: first vertex address (nmesh x 1) + mesh_faceadr: first face address (nmesh x 1) + mesh_vert: vertex positions for all meshes (nmeshvert, 3) + mesh_face: vertex face data (nmeshface, 3) geom_convex_face: vertex face data, MJX only (ngeom,) geom_convex_vert: vertex data, MJX only (ngeom,) geom_convex_edge: unique edge data, MJX only (ngeom,) @@ -390,6 +397,8 @@ class Model(PyTreeNode): ngeom: int nsite: int nmesh: int + nmeshvert: int + nmeshface: int nmat: int npair: int nexclude: int @@ -447,6 +456,7 @@ class Model(PyTreeNode): geom_conaffinity: np.ndarray geom_condim: np.ndarray geom_bodyid: np.ndarray + geom_dataid: np.ndarray geom_group: np.ndarray geom_matid: np.ndarray geom_priority: np.ndarray @@ -463,6 +473,10 @@ class Model(PyTreeNode): site_bodyid: np.ndarray site_pos: jax.Array site_quat: jax.Array + mesh_vertadr: np.ndarray + mesh_faceadr: np.ndarray + mesh_vert: np.ndarray + mesh_face: np.ndarray mat_rgba: np.ndarray pair_dim: np.ndarray pair_geom1: np.ndarray diff --git a/mjx/mujoco/mjx/test_data/ray.xml b/mjx/mujoco/mjx/test_data/ray.xml index a6424ec4..bbf42b1e 100644 --- a/mjx/mujoco/mjx/test_data/ray.xml +++ b/mjx/mujoco/mjx/test_data/ray.xml @@ -1,6 +1,7 @@ + @@ -12,5 +13,6 @@ + From 71b056175843065fae32892681572ed1a9cdf677 Mon Sep 17 00:00:00 2001 From: Kyle Bayes Date: Wed, 17 Jan 2024 04:20:24 -0800 Subject: [PATCH 13/92] Clean up formatting issues in user_objects.h PiperOrigin-RevId: 599127766 Change-Id: If752c00038818fb21476eb98be891ead467aad56 --- src/user/user_objects.h | 120 +++++++++++++++++++++------------------- 1 file changed, 63 insertions(+), 57 deletions(-) diff --git a/src/user/user_objects.h b/src/user/user_objects.h index 61687868..c4839fee 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -16,12 +16,17 @@ #define MUJOCO_SRC_USER_USER_OBJECTS_H_ #include +#include #include #include #include +#include +#include #include #include "lodepng.h" + +#include #include #include @@ -331,37 +336,38 @@ class mjCJoint : public mjCBase { public: // variables set by user: joint properties - mjtJoint type; // type of Joint - int group; // used for rendering - int limited; // does joint have limits: 0 false, 1 true, 2 auto - int actfrclimited; // are actuator forces on joints limited: 0 false, 1 true, 2 auto - double pos[3]; // anchor position - double axis[3]; // joint axis - double stiffness; // stiffness coefficient - double springdamper[2]; // timeconst, dampratio - double range[2]; // joint limits - double actfrcrange[2]; // actuator force limits - mjtNum solref_limit[mjNREF]; // solver reference: joint limits - mjtNum solimp_limit[mjNIMP]; // solver impedance: joint limits - mjtNum solref_friction[mjNREF]; // solver reference: dof friction - mjtNum solimp_friction[mjNIMP]; // solver impedance: dof friction - double margin; // margin value for joint limit detection - double ref; // value at reference configuration: qpos0 - double springref; // spring reference value: qpos_spring - std::vector userdata; // user data + mjtJoint type; // type of Joint + int group; // used for rendering + int limited; // does joint have limits: 0 false, 1 true, 2 auto + int actfrclimited; // are actuator forces on joints limited: 0 false, 1 true, 2 auto + double pos[3]; // anchor position + double axis[3]; // joint axis + double stiffness; // stiffness coefficient + double springdamper[2]; // timeconst, dampratio + double range[2]; // joint limits + double actfrcrange[2]; // actuator force limits + mjtNum solref_limit[mjNREF]; // solver reference: joint limits + mjtNum solimp_limit[mjNIMP]; // solver impedance: joint limits + mjtNum solref_friction[mjNREF]; // solver reference: dof friction + mjtNum solimp_friction[mjNIMP]; // solver impedance: dof friction + double margin; // margin value for joint limit detection + double ref; // value at reference configuration: qpos0 + double springref; // spring reference value: qpos_spring + std::vector userdata; // user data // variables set by user: dof properties - double armature; // armature inertia (mass for slider) - double damping; // damping coefficient - double frictionloss; // friction loss + double armature; // armature inertia (mass for slider) + double damping; // damping coefficient + double frictionloss; // friction loss - double urdfeffort; // store effort field from urdf + double urdfeffort; // store effort field from urdf private: - mjCJoint(mjCModel* = 0, mjCDef* = 0);// constructor - int Compile(void); // compiler; return dofnum + mjCJoint(mjCModel* = 0, mjCDef* = 0); - mjCBody* body; // joint's body + int Compile(void); // compiler; return dofnum + + mjCBody* body; // joint's body }; @@ -429,19 +435,19 @@ class mjCGeom : public mjCBase { double quat[4]; // orientation private: - mjCGeom(mjCModel* = 0, mjCDef* = 0);// constructor + mjCGeom(mjCModel* = 0, mjCDef* = 0); void Compile(void); // compiler double GetRBound(void); // compute bounding sphere radius void ComputeAABB(void); // compute axis-aligned bounding box - bool visual_; // true: geom does not collide and is unreferenced - int matid; // id of geom's material - mjCMesh* mesh; // geom's mesh - mjCHField* hfield; // geom's hfield - double mass; // mass - double inertia[3]; // local diagonal inertia - double aabb[6]; // axis-aligned bounding box (center, size) - mjCBody* body; // geom's body + bool visual_; // true: geom does not collide and is unreferenced + int matid; // id of geom's material + mjCMesh* mesh; // geom's mesh + mjCHField* hfield; // geom's hfield + double mass; // mass + double inertia[3]; // local diagonal inertia + double aabb[6]; // axis-aligned bounding box (center, size) + mjCBody* body; // geom's body }; @@ -598,24 +604,24 @@ class mjCFlex: public mjCBase { void DelTexcoord(); // delete texcoord private: - mjCFlex(mjCModel* = 0); // constructor - void Compile(const mjVFS* vfs); // compiler - void CreateBVH(void); // create flex BVH - void CreateShellPair(void); // create shells and evpairs + mjCFlex(mjCModel* = 0); + void Compile(const mjVFS* vfs); // compiler + void CreateBVH(void); // create flex BVH + void CreateShellPair(void); // create shells and evpairs - int nvert; // number of verices - int nedge; // number of edges - int nelem; // number of elements - int matid; // material id - bool rigid; // all vertices attached to the same body - bool centered; // all vertices coordinates (0,0,0) - std::vector vertbodyid; // vertex body ids - std::vector> edge; // edge vertex ids - std::vector shell; // shell fragment vertex ids (dim per fragment) - std::vector elemlayer; // element layer (distance from border) - std::vector evpair; // element-vertex pairs - std::vector vertxpos; // global vertex positions - mjCBoundingVolumeHierarchy tree; // bounding volume hierarchy + int nvert; // number of verices + int nedge; // number of edges + int nelem; // number of elements + int matid; // material id + bool rigid; // all vertices attached to the same body + bool centered; // all vertices coordinates (0,0,0) + std::vector vertbodyid; // vertex body ids + std::vector> edge; // edge vertex ids + std::vector shell; // shell fragment vertex ids (dim per fragment) + std::vector elemlayer; // element layer (distance from border) + std::vector evpair; // element-vertex pairs + std::vector vertxpos; // global vertex positions + mjCBoundingVolumeHierarchy tree; // bounding volume hierarchy }; @@ -624,7 +630,7 @@ class mjCFlex: public mjCBase { // Describes a mesh class mjCMesh: public mjCBase { - friend class mjCFlexcomp; + friend class mjCFlexcomp; public: mjCMesh(mjCModel* = 0, mjCDef* = 0); ~mjCMesh(); @@ -644,7 +650,7 @@ class mjCMesh: public mjCBase { const std::vector& usertexcoord() const { return usertexcoord_; } const std::vector& userface() const { return userface_; } - // mesh properites computed by Compile + // mesh properties computed by Compile const double* aamm() const { return aamm_; } // number of vertices, normals, texture coordinates, and faces @@ -892,11 +898,11 @@ class mjCTexture : public mjCBase { void Builtin2D(void); // make builtin 2D void BuiltinCube(void); // make builtin cube - void Load2D(std::string filename, const mjVFS* vfs); // load 2D from file - void LoadCubeSingle(std::string filename, const mjVFS* vfs); // load cube from single file - void LoadCubeSeparate(const mjVFS* vfs); // load cube from separate files + void Load2D(std::string filename, const mjVFS* vfs); // load 2D from file + void LoadCubeSingle(std::string filename, const mjVFS* vfs); // load cube from single file + void LoadCubeSeparate(const mjVFS* vfs); // load cube from separate files - void LoadFlip(std::string filename, const mjVFS* vfs, // load and flip + void LoadFlip(std::string filename, const mjVFS* vfs, // load and flip std::vector& image, unsigned int& w, unsigned int& h); From 67a00a3a510ca53e6699bb9121d4c671b9e29e96 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Wed, 17 Jan 2024 04:33:00 -0800 Subject: [PATCH 14/92] Rename `mjMAXTHREADS` -> `mjMAXTHREAD` (consistency). PiperOrigin-RevId: 599130235 Change-Id: Ia80379b8ae0574a7cd8fbc668657f806e90d6c5c --- doc/APIreference/APIglobals.rst | 2 +- doc/includes/references.h | 2 +- include/mujoco/mjdata.h | 2 +- include/mujoco/mjthread.h | 2 +- include/mujoco/mjxmacro.h | 2 +- src/engine/engine_io.c | 2 +- src/thread/thread_pool.cc | 2 +- unity/Runtime/Bindings/MjBindings.cs | 2 +- 8 files changed, 8 insertions(+), 8 deletions(-) diff --git a/doc/APIreference/APIglobals.rst b/doc/APIreference/APIglobals.rst index 483383c4..7e3778a5 100644 --- a/doc/APIreference/APIglobals.rst +++ b/doc/APIreference/APIglobals.rst @@ -489,7 +489,7 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr - 1000 - Maximum number of textures allowed. Defined in `mjrender.h `_. - * - ``mjMAXTHREADS`` + * - ``mjMAXTHREAD`` - 128 - Maximum number OS threads that can be used in a thread pool. Defined in `mjthread.h `_. diff --git a/doc/includes/references.h b/doc/includes/references.h index 1c788150..cfee19fb 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -148,7 +148,7 @@ struct mjData_ { // memory utilization stats size_t maxuse_stack; // maximum stack allocation in bytes - size_t maxuse_threadstack[mjMAXTHREADS]; // maximum stack allocation per thread in bytes + size_t maxuse_threadstack[mjMAXTHREAD]; // maximum stack allocation per thread in bytes size_t maxuse_arena; // maximum arena allocation in bytes int maxuse_con; // maximum number of contacts int maxuse_efc; // maximum number of scalar constraints diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index 1fc6fe19..53a42498 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -176,7 +176,7 @@ struct mjData_ { // memory utilization stats size_t maxuse_stack; // maximum stack allocation in bytes - size_t maxuse_threadstack[mjMAXTHREADS]; // maximum stack allocation per thread in bytes + size_t maxuse_threadstack[mjMAXTHREAD]; // maximum stack allocation per thread in bytes size_t maxuse_arena; // maximum arena allocation in bytes int maxuse_con; // maximum number of contacts int maxuse_efc; // maximum number of scalar constraints diff --git a/include/mujoco/mjthread.h b/include/mujoco/mjthread.h index 9a526fa5..153ff0cc 100644 --- a/include/mujoco/mjthread.h +++ b/include/mujoco/mjthread.h @@ -15,7 +15,7 @@ #ifndef MUJOCO_INCLUDE_MJTHREAD_H_ #define MUJOCO_INCLUDE_MJTHREAD_H_ -#define mjMAXTHREADS 128 // maximum number of threads in a thread pool +#define mjMAXTHREAD 128 // maximum number of threads in a thread pool typedef enum mjtTaskStatus_ { // status values for mjTask mjTASK_NEW = 0, // newly created diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 86f61e0c..a46aea71 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -739,7 +739,7 @@ // vector fields of mjData #define MJDATA_VECTOR \ - X( size_t, maxuse_threadstack, mjMAXTHREADS, 1 ) \ + X( size_t, maxuse_threadstack, mjMAXTHREAD, 1 ) \ X( mjWarningStat, warning, mjNWARNING, 1 ) \ X( mjTimerStat, timer, mjNTIMER, 1 ) \ X( mjSolverStat, solver, mjNILSAND, mjNSOLVER ) \ diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index ca55358a..7c44b54f 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -1568,7 +1568,7 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) { // clear memory utilization stats d->maxuse_stack = 0; - mju_zeroSizeT(d->maxuse_threadstack, mjMAXTHREADS); + mju_zeroSizeT(d->maxuse_threadstack, mjMAXTHREAD); d->maxuse_arena = 0; d->maxuse_con = 0; d->maxuse_efc = 0; diff --git a/src/thread/thread_pool.cc b/src/thread/thread_pool.cc index 6c32c2b8..9bb03133 100644 --- a/src/thread/thread_pool.cc +++ b/src/thread/thread_pool.cc @@ -65,7 +65,7 @@ class ThreadPoolImpl : public mjThreadPool { public: ThreadPoolImpl(int num_worker) : mjThreadPool{num_worker} { // initialize worker threads - for (int i = 0; i < std::min(num_worker, mjMAXTHREADS); ++i) { + for (int i = 0; i < std::min(num_worker, mjMAXTHREAD); ++i) { WorkerThread worker{ std::make_unique(ThreadPoolWorker, this, i)}; workers_.push_back(std::move(worker)); diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index fd2c33bf..1aeb7a12 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -58,7 +58,7 @@ public const bool THIRD_PARTY_MUJOCO_MJRENDER_H_ = true; public const int mjNAUX = 10; public const int mjMAXTEXTURE = 1000; public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJTHREAD_H_ = true; -public const int mjMAXTHREADS = 128; +public const int mjMAXTHREAD = 128; public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJTNUM_H_ = true; public const bool mjUSEDOUBLE = true; public const double mjMINVAL = 1e-15; From a14a584f1d506c8636342af6f39e5e7157966a1a Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Wed, 17 Jan 2024 07:19:11 -0800 Subject: [PATCH 15/92] Fix bug in muscle length-gain curve. Fixes #1342 PiperOrigin-RevId: 599165049 Change-Id: Ic5ef77b4349a8a9c343eacfd781195bfebea9aca --- doc/changelog.rst | 2 ++ src/engine/engine_derivative.c | 20 +----------- src/engine/engine_util_misc.c | 47 +++++++++++++++++----------- src/engine/engine_util_misc.h | 3 ++ test/engine/engine_util_misc_test.cc | 13 ++++++++ 5 files changed, 47 insertions(+), 38 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index eb9aeab0..1f97f21f 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -21,6 +21,8 @@ Bug fixes ^^^^^^^^^ 6. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes :github:issue:`1270`. +7. Fixed a bug in the :ref:`muscle model` that led to non-zero values outside the lower + bound of the length range. Fixes :github:issue:`1342`. Version 3.1.1 (December 18, 2023) diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c index c0c0d5fe..643c9b4f 100644 --- a/src/engine/engine_derivative.c +++ b/src/engine/engine_derivative.c @@ -790,11 +790,6 @@ static mjtNum mjd_muscleGain_vel(mjtNum len, mjtNum vel, const mjtNum lengthrang force = scale / mju_max(mjMINVAL, acc0); } - // mid-ranges - mjtNum a = 0.5*(lmin+1); - mjtNum b = 0.5*(1+lmax); - mjtNum x; - // optimum length mjtNum L0 = (lengthrange[1]-lengthrange[0]) / mju_max(mjMINVAL, range[1]-range[0]); @@ -803,20 +798,7 @@ static mjtNum mjd_muscleGain_vel(mjtNum len, mjtNum vel, const mjtNum lengthrang mjtNum V = vel / mju_max(mjMINVAL, L0*vmax); // length curve - mjtNum FL = 0; - if (L >= lmin && L <= a) { - x = (L-lmin) / mju_max(mjMINVAL, a-lmin); - FL = 0.5*x*x; - } else if (L <= 1) { - x = (1-L) / mju_max(mjMINVAL, 1-a); - FL = 1 - 0.5*x*x; - } else if (L <= b) { - x = (L-1) / mju_max(mjMINVAL, b-1); - FL = 1 - 0.5*x*x; - } else if (L <= lmax) { - x = (lmax-L) / mju_max(mjMINVAL, lmax-b); - FL = 0.5*x*x; - } + mjtNum FL = mju_muscleGainLength(L, lmin, lmax); // velocity curve mjtNum dFV; diff --git a/src/engine/engine_util_misc.c b/src/engine/engine_util_misc.c index 37eca9e4..876fe52a 100644 --- a/src/engine/engine_util_misc.c +++ b/src/engine/engine_util_misc.c @@ -455,6 +455,33 @@ void mju_geomSemiAxes(const mjModel* m, int geom_id, mjtNum semiaxes[3]) { //------------------------------ actuator models --------------------------------------------------- +// normalized muscle length-gain curve +mjtNum mju_muscleGainLength(mjtNum length, mjtNum lmin, mjtNum lmax) { + if (lmin <= length && length <= lmax) { + // mid-ranges (maximum is at 1.0) + mjtNum a = 0.5*(lmin+1); + mjtNum b = 0.5*(1+lmax); + + if (length <= a) { + mjtNum x = (length-lmin) / mjMAX(mjMINVAL, a-lmin); + return 0.5*x*x; + } else if (length <= 1) { + mjtNum x = (1-length) / mjMAX(mjMINVAL, 1-a); + return 1 - 0.5*x*x; + } else if (length <= b) { + mjtNum x = (length-1) / mjMAX(mjMINVAL, b-1); + return 1 - 0.5*x*x; + } else { + mjtNum x = (lmax-length) / mjMAX(mjMINVAL, lmax-b); + return 0.5*x*x; + } + } + + return 0.0; +} + + + // muscle active force, prm = (range[2], force, scale, lmin, lmax, vmax, fpmax, fvmax) mjtNum mju_muscleGain(mjtNum len, mjtNum vel, const mjtNum lengthrange[2], mjtNum acc0, const mjtNum prm[9]) { @@ -472,11 +499,6 @@ mjtNum mju_muscleGain(mjtNum len, mjtNum vel, const mjtNum lengthrange[2], force = scale / mjMAX(mjMINVAL, acc0); } - // mid-ranges - mjtNum a = 0.5*(lmin+1); - mjtNum b = 0.5*(1+lmax); - mjtNum x; - // optimum length mjtNum L0 = (lengthrange[1]-lengthrange[0]) / mjMAX(mjMINVAL, range[1]-range[0]); @@ -485,20 +507,7 @@ mjtNum mju_muscleGain(mjtNum len, mjtNum vel, const mjtNum lengthrange[2], mjtNum V = vel / mjMAX(mjMINVAL, L0*vmax); // length curve - mjtNum FL = 0; - if (L >= lmin && L <= a) { - x = (L-lmin) / mjMAX(mjMINVAL, a-lmin); - FL = 0.5*x*x; - } else if (L <= 1) { - x = (1-L) / mjMAX(mjMINVAL, 1-a); - FL = 1 - 0.5*x*x; - } else if (L <= b) { - x = (L-1) / mjMAX(mjMINVAL, b-1); - FL = 1 - 0.5*x*x; - } else if (L <= lmax) { - x = (lmax-L) / mjMAX(mjMINVAL, lmax-b); - FL = 0.5*x*x; - } + mjtNum FL = mju_muscleGainLength(L, lmin, lmax); // velocity curve mjtNum FV; diff --git a/src/engine/engine_util_misc.h b/src/engine/engine_util_misc.h index cb317258..21bc3f83 100644 --- a/src/engine/engine_util_misc.h +++ b/src/engine/engine_util_misc.h @@ -33,6 +33,9 @@ mjtNum mju_wrap(mjtNum* wpnt, const mjtNum* x0, const mjtNum* x1, const mjtNum* xpos, const mjtNum* xmat, const mjtNum* size, int type, const mjtNum* side); +// normalized muscle length-gain curve +MJAPI mjtNum mju_muscleGainLength(mjtNum length, mjtNum lmin, mjtNum lmax); + // muscle active force, prm = (range[2], force, scale, lmin, lmax, vmax, fpmax, fvmax) MJAPI mjtNum mju_muscleGain(mjtNum len, mjtNum vel, const mjtNum lengthrange[2], mjtNum acc0, const mjtNum prm[9]); diff --git a/test/engine/engine_util_misc_test.cc b/test/engine/engine_util_misc_test.cc index 71702b8c..97557914 100644 --- a/test/engine/engine_util_misc_test.cc +++ b/test/engine/engine_util_misc_test.cc @@ -135,6 +135,19 @@ TEST_F(MujocoTest, SmoothMuscleDynamics) { } } +TEST_F(MujocoTest, MuscleGainLength) { + mjtNum lmin = 0.5; + mjtNum lmax = 1.5; + + EXPECT_EQ(mju_muscleGainLength(0.0, lmin, lmax), 0); + EXPECT_EQ(mju_muscleGainLength(0.5, lmin, lmax), 0); + EXPECT_EQ(mju_muscleGainLength(0.75, lmin, lmax), 0.5); + EXPECT_EQ(mju_muscleGainLength(1.0, lmin, lmax), 1); + EXPECT_EQ(mju_muscleGainLength(1.25, lmin, lmax), 0.5); + EXPECT_EQ(mju_muscleGainLength(1.5, lmin, lmax), 0); + EXPECT_EQ(mju_muscleGainLength(2.0, lmin, lmax), 0); +} + TEST_F(MujocoTest, mju_makefullname) { char buffer[1000]; constexpr char path[] = "engine/testdata/"; From edc254b213e21b52391acf9ad4dfe1951c0c011e Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Wed, 17 Jan 2024 10:37:50 -0800 Subject: [PATCH 16/92] Update MJX doc for #1344. PiperOrigin-RevId: 599220357 Change-Id: I532f42ed086da59632eeaac0dcd3de6ffcf93cd4 --- doc/mjx.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/mjx.rst b/doc/mjx.rst index aafb7a86..251e05fc 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -230,7 +230,7 @@ The following features are **in development** and coming soon: * - :ref:`Geom ` - ``HFIELD``, ``ELLIPSOID``, ``CYLINDER`` * - :ref:`Constraint ` - - ``CONTACT_FRICTIONLESS``, ``CONTACT_ELLIPTIC``, ``FRICTION_DOF`` + - :ref:`Frictionloss `, ``CONTACT_FRICTIONLESS``, ``CONTACT_ELLIPTIC``, ``FRICTION_DOF`` * - :ref:`Integrator ` - ``IMPLICIT``, ``IMPLICITFAST`` * - :ref:`Cone ` From 2feefbc5d24899ea3f6ebe1e78e57555d0136db2 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Wed, 17 Jan 2024 10:53:19 -0800 Subject: [PATCH 17/92] Remove fine-grained timing from narrow-phase collision functions, delete `mjTIMER_COL_MID`. PiperOrigin-RevId: 599225419 Change-Id: Ieca166753c20a0abf660e2821b2844dda8758951 --- doc/changelog.rst | 13 +++++++----- doc/includes/references.h | 1 - include/mujoco/mjdata.h | 1 - introspect/enums.py | 5 ++--- introspect/structs.py | 2 +- sample/testspeed.cc | 2 +- src/engine/engine_collision_driver.c | 30 +++------------------------- src/engine/engine_support.c | 1 - unity/Runtime/Bindings/MjBindings.cs | 6 ++---- 9 files changed, 17 insertions(+), 44 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 1f97f21f..33230ae5 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -9,17 +9,20 @@ General ^^^^^^^ 1. Improved the :ref:discardvisual compiler flag, which now discards all visual-only assets. See :ref:discardvisual for details. +2. Removed the :ref:`timer` for midphase colllision detection, it is now folded in with the narrowphase + timer. This is because timing the two phases seperately required fine-grained timers inside the collision + functions; these functions are so small and fast that the timer itself was incurring a measurable cost. MJX ^^^ -2. Added :ref:`dyntype` ``filterexact``. -3. Added :at:`site` transmission. -4. Updated MJX colab tutorial with more stable quadruped environment. -5. Added ``mjx.ray`` which mirrors :ref:`mj_ray` for planes, spheres, capsules, boxes, and meshes. +3. Added :ref:`dyntype` ``filterexact``. +4. Added :at:`site` transmission. +5. Updated MJX colab tutorial with more stable quadruped environment. +6. Added ``mjx.ray`` which mirrors :ref:`mj_ray` for planes, spheres, capsules, boxes, and meshes. Bug fixes ^^^^^^^^^ -6. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes +7. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes :github:issue:`1270`. 7. Fixed a bug in the :ref:`muscle model` that led to non-zero values outside the lower bound of the length range. Fixes :github:issue:`1342`. diff --git a/doc/includes/references.h b/doc/includes/references.h index cfee19fb..01504204 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -75,7 +75,6 @@ typedef enum mjtTimer_ { // internal timers // breakdown of mj_collision mjTIMER_COL_BROAD, // broadphase - mjTIMER_COL_MID, // midphase mjTIMER_COL_NARROW, // narrowphase mjNTIMER // number of timers diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index 53a42498..65886b38 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -87,7 +87,6 @@ typedef enum mjtTimer_ { // internal timers // breakdown of mj_collision mjTIMER_COL_BROAD, // broadphase - mjTIMER_COL_MID, // midphase mjTIMER_COL_NARROW, // narrowphase mjNTIMER // number of timers diff --git a/introspect/enums.py b/introspect/enums.py index 8502dbdb..180c157b 100644 --- a/introspect/enums.py +++ b/introspect/enums.py @@ -458,9 +458,8 @@ ENUMS: Mapping[str, EnumDecl] = dict([ ('mjTIMER_POS_MAKE', 11), ('mjTIMER_POS_PROJECT', 12), ('mjTIMER_COL_BROAD', 13), - ('mjTIMER_COL_MID', 14), - ('mjTIMER_COL_NARROW', 15), - ('mjNTIMER', 16), + ('mjTIMER_COL_NARROW', 14), + ('mjNTIMER', 15), ]), )), ('mjtCatBit', diff --git a/introspect/structs.py b/introspect/structs.py index 43e7cdc9..02e92c7f 100644 --- a/introspect/structs.py +++ b/introspect/structs.py @@ -4154,7 +4154,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([ name='timer', type=ArrayType( inner_type=ValueType(name='mjTimerStat'), - extents=(16,), + extents=(15,), ), doc='timer statistics', ), diff --git a/sample/testspeed.cc b/sample/testspeed.cc index 03d67d15..66dabf3a 100644 --- a/sample/testspeed.cc +++ b/sample/testspeed.cc @@ -278,7 +278,7 @@ int main(int argc, char** argv) { // components of mjTIMER_POS_COLLISION if (i == mjTIMER_POS_COLLISION) { - for (int j : {mjTIMER_COL_BROAD, mjTIMER_COL_MID, mjTIMER_COL_NARROW}) { + for (int j : {mjTIMER_COL_BROAD, mjTIMER_COL_NARROW}) { int number = d[0]->timer[j].number; mjtNum jstep = number ? d[0]->timer[j].duration/number : 0.0; mjtNum percent = number ? 100*jstep/tstep : 0.0; diff --git a/src/engine/engine_collision_driver.c b/src/engine/engine_collision_driver.c index f6132b74..b8a8a7d0 100644 --- a/src/engine/engine_collision_driver.c +++ b/src/engine/engine_collision_driver.c @@ -291,12 +291,9 @@ void mj_collision(const mjModel* m, mjData* d) { unsigned int last_signature = -1; TM_END(mjTIMER_COL_BROAD); - // midphase collision detector + // narrowphase and midphase collision detector TM_RESTART; - // save current narrowphase duration - mjtNum tmNarrow = d->timer[mjTIMER_COL_NARROW].duration; - // process bodyflex pairs returned by broadphase, merge with predefined geom pairs int pairadr = 0; for (int i=0; i < nbfpair; i++) { @@ -471,14 +468,8 @@ void mj_collision(const mjModel* m, mjData* d) { } } - // end midphase timer - TM_END(mjTIMER_COL_MID); - - // subtract nested narrowphase timing from midphase timer - d->timer[mjTIMER_COL_MID].duration -= (d->timer[mjTIMER_COL_NARROW].duration - tmNarrow); - - // increment narrowphase counter - d->timer[mjTIMER_COL_NARROW].number++; + // end narrowphase and midphase timer + TM_END(mjTIMER_COL_NARROW); mj_freeStack(d); TM_END1(mjTIMER_POS_COLLISION); @@ -1444,8 +1435,6 @@ static void mj_makeCapsule(const mjModel* m, mjData* d, int f, const int vid[2], // test two geoms for collision, apply filters, add to contact list void mj_collideGeoms(const mjModel* m, mjData* d, int g1, int g2) { - TM_START; - int num, type1, type2, condim; mjtNum margin, gap, friction[5], solref[mjNREF], solimp[mjNIMP]; mjtNum solreffriction[mjNREF] = {0}; @@ -1635,9 +1624,6 @@ void mj_collideGeoms(const mjModel* m, mjData* d, int g1, int g2) { // move arena pointer back to the end of the contact array resetArena(d); - - // add duration without incrementing counter - TM_ADD(mjTIMER_COL_NARROW); } @@ -1850,8 +1836,6 @@ void mj_collideFlexSAP(const mjModel* m, mjData* d, int f) { // test a geom and an elem for collision, add to contact list void mj_collideGeomElem(const mjModel* m, mjData* d, int g, int f, int e) { - TM_START; - mjtNum margin = mj_assignMargin(m, mju_max(m->geom_margin[g], m->flex_margin[f])); int dim = m->flex_dim[f], type = m->geom_type[g]; int num; @@ -1965,17 +1949,12 @@ void mj_collideGeomElem(const mjModel* m, mjData* d, int g, int f, int e) { // move arena pointer back to the end of the contact array resetArena(d); - - // add duration without incrementing counter - TM_ADD(mjTIMER_COL_NARROW); } // test two elems for collision, add to contact list void mj_collideElems(const mjModel* m, mjData* d, int f1, int e1, int f2, int e2) { - TM_START; - mjtNum margin = mj_assignMargin(m, mju_max(m->flex_margin[f1], m->flex_margin[f2])); int dim1 = m->flex_dim[f1], dim2 = m->flex_dim[f2]; int num; @@ -2070,9 +2049,6 @@ void mj_collideElems(const mjModel* m, mjData* d, int f1, int e1, int f2, int e2 // move arena pointer back to the end of the contact array resetArena(d); - - // add duration without incrementing counter - TM_ADD(mjTIMER_COL_NARROW); } diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index dfa91d0a..f44baaf7 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -89,7 +89,6 @@ const char* mjTIMERSTRING[mjNTIMER]= { "pos_make", "pos_project", "col_broadphase", - "col_midphase", "col_narrowphase" }; diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 1aeb7a12..ad6bc8b9 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -138,9 +138,8 @@ public enum mjtTimer : int{ mjTIMER_POS_MAKE = 11, mjTIMER_POS_PROJECT = 12, mjTIMER_COL_BROAD = 13, - mjTIMER_COL_MID = 14, - mjTIMER_COL_NARROW = 15, - mjNTIMER = 16, + mjTIMER_COL_NARROW = 14, + mjNTIMER = 15, } public enum mjtDisableBit : int{ mjDSBL_CONSTRAINT = 1, @@ -769,7 +768,6 @@ public unsafe struct mjData_ { public mjTimerStat_ timer12; public mjTimerStat_ timer13; public mjTimerStat_ timer14; - public mjTimerStat_ timer15; public mjSolverStat_ solver0; public mjSolverStat_ solver1; public mjSolverStat_ solver2; From fea7c10b9568fd470e8d03b27ceb7efeb8cb4563 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Fri, 19 Jan 2024 02:59:06 -0800 Subject: [PATCH 18/92] Add `bvhactive` flag to `visual/global`. Fixes #1279. PiperOrigin-RevId: 599778272 Change-Id: If5530629035cc68b198447503388f59a04041b13 --- doc/XMLreference.rst | 15 ++++++++ doc/XMLschema.rst | 6 ++-- doc/changelog.rst | 21 ++++++++---- doc/includes/references.h | 3 ++ include/mujoco/mjmodel.h | 3 ++ introspect/structs.py | 21 ++++++++++++ python/mujoco/structs.cc | 3 ++ python/mujoco/structs.h | 2 ++ simulate/simulate.cc | 1 + src/engine/engine_collision_driver.c | 29 +++++++++++----- src/engine/engine_collision_sdf.c | 9 ++--- src/engine/engine_io.c | 3 ++ src/engine/engine_ray.c | 8 ++++- src/engine/engine_vis_visualize.c | 38 +++++++++++++-------- src/xml/xml_native_reader.cc | 14 +++++--- src/xml/xml_native_writer.cc | 3 ++ test/engine/testdata/ray/stanford_bunny.xml | 9 ++++- unity/Runtime/Bindings/MjBindings.cs | 3 ++ 18 files changed, 147 insertions(+), 44 deletions(-) diff --git a/doc/XMLreference.rst b/doc/XMLreference.rst index 072aaab5..c781f400 100644 --- a/doc/XMLreference.rst +++ b/doc/XMLreference.rst @@ -715,6 +715,12 @@ is effectively a miscellaneous subsection. This attribute specifies how the equivalent inertia is visualized. "false": use box, "true": use ellipsoid. +.. _visual-global-bvactive: + +:at:`bvactive`: :at-val:`[false, true], "true"` + This attribute specifies whether collision and raycasting code should mark elements of Bounding Volume Hierarchies + as intersecting, for the purpose of visualization. Setting this attribute to "false" can speed up simulation for + models with high-resolution meshes. .. _visual-quality: @@ -1122,6 +1128,15 @@ disables the rendering of the corresponding object. :at:`frustum`: :at-val:`real(4), "1 1 0 0.2"` Color used to render the camera frustum. +.. _visual-rgba-bv: + +:at:`bv`: :at-val:`real(4), "0 1 0 0.5"` + Color used to render bounding volumes. + +.. _visual-rgba-bvactive: + +:at:`bvactive`: :at-val:`real(4), "1 0 0 0.5"` + Color used to render active bounding volumes, if the :ref:`bvactive` flag is "true". .. _asset: diff --git a/doc/XMLschema.rst b/doc/XMLschema.rst index 702adeb2..b319f26b 100644 --- a/doc/XMLschema.rst +++ b/doc/XMLschema.rst @@ -65,7 +65,7 @@ | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`linewidth` | :ref:`glow` | :ref:`offwidth` | :ref:`offheight` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`realtime` | :ref:`ellipsoidinertia` | | | | +| | | | :ref:`realtime` | :ref:`ellipsoidinertia` | :ref:`bvactive` | | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_| visual |br| |_| |L| | | .. table:: | @@ -126,7 +126,9 @@ | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`contacttorque` | :ref:`contactgap` | :ref:`rangefinder` | :ref:`constraint` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`slidercrank` | :ref:`crankbroken` | :ref:`frustum` | | | +| | | | :ref:`slidercrank` | :ref:`crankbroken` | :ref:`frustum` | :ref:`bv` | | +| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +| | | | :ref:`bvactive` | | | | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | mujoco |br| |L| | | *no attributes* | diff --git a/doc/changelog.rst b/doc/changelog.rst index 33230ae5..910d6048 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -12,19 +12,26 @@ General 2. Removed the :ref:`timer` for midphase colllision detection, it is now folded in with the narrowphase timer. This is because timing the two phases seperately required fine-grained timers inside the collision functions; these functions are so small and fast that the timer itself was incurring a measurable cost. +3. Added the flag :ref:`bvactive` to ``visual/global``, allowing users to turn off + visualisation of active bounding volumes (the red/green boxes in this :ref:`this changelog item`). For + models with very high-resolution meshes, the computation required for this visualization can slow down simulation + speed. Fixes :github:issue:`1279`. + + - Added color of :ref:`bounding volumes` and :ref:`active bounding volumes` + to :ref:`visual/rgba`. MJX ^^^ -3. Added :ref:`dyntype` ``filterexact``. -4. Added :at:`site` transmission. -5. Updated MJX colab tutorial with more stable quadruped environment. -6. Added ``mjx.ray`` which mirrors :ref:`mj_ray` for planes, spheres, capsules, boxes, and meshes. +4. Added :ref:`dyntype` ``filterexact``. +5. Added :at:`site` transmission. +6. Updated MJX colab tutorial with more stable quadruped environment. +7. Added ``mjx.ray`` which mirrors :ref:`mj_ray` for planes, spheres, capsules, boxes, and meshes. Bug fixes ^^^^^^^^^ -7. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes +8. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes :github:issue:`1270`. -7. Fixed a bug in the :ref:`muscle model` that led to non-zero values outside the lower +9. Fixed a bug in the :ref:`muscle model` that led to non-zero values outside the lower bound of the length range. Fixes :github:issue:`1342`. @@ -624,6 +631,8 @@ General :align: right :width: 350px +.. _midphase: + 2. Added a collision mid-phase for pruning geoms in body pairs, see :ref:`documentation` for more details. This is based on static AABB bounding volume hierarchy (a BVH binary tree) in the body inertial frame. The GIF on the right is cut from `this longer video `__. diff --git a/doc/includes/references.h b/doc/includes/references.h index 01504204..b94d5d71 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -753,6 +753,7 @@ struct mjVisual_ { // visualization options int offwidth; // width of offscreen buffer int offheight; // height of offscreen buffer int ellipsoidinertia; // geom for inertia visualization (0: box, 1: ellipsoid) + int bvactive; // visualize active bounding volumes (0: no, 1: yes) } global; struct { // rendering quality @@ -830,6 +831,8 @@ struct mjVisual_ { // visualization options float slidercrank[4]; // slidercrank float crankbroken[4]; // used when crank must be stretched/broken float frustum[4]; // camera frustum + float bv[4]; // bounding volume + float bvactive[4]; // active bounding volume } rgba; }; typedef struct mjVisual_ mjVisual; diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index d3417683..b91ee594 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -459,6 +459,7 @@ struct mjVisual_ { // visualization options int offwidth; // width of offscreen buffer int offheight; // height of offscreen buffer int ellipsoidinertia; // geom for inertia visualization (0: box, 1: ellipsoid) + int bvactive; // visualize active bounding volumes (0: no, 1: yes) } global; struct { // rendering quality @@ -536,6 +537,8 @@ struct mjVisual_ { // visualization options float slidercrank[4]; // slidercrank float crankbroken[4]; // used when crank must be stretched/broken float frustum[4]; // camera frustum + float bv[4]; // bounding volume + float bvactive[4]; // active bounding volume } rgba; }; typedef struct mjVisual_ mjVisual; diff --git a/introspect/structs.py b/introspect/structs.py index 02e92c7f..8f44927d 100644 --- a/introspect/structs.py +++ b/introspect/structs.py @@ -352,6 +352,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=ValueType(name='int'), doc='geom for inertia visualization (0: box, 1: ellipsoid)', # pylint: disable=line-too-long ), + StructFieldDecl( + name='bvactive', + type=ValueType(name='int'), + doc='visualize active bounding volumes (0: no, 1: yes)', # pylint: disable=line-too-long + ), ), ), doc='', @@ -780,6 +785,22 @@ STRUCTS: Mapping[str, StructDecl] = dict([ ), doc='camera frustum', ), + StructFieldDecl( + name='bv', + type=ArrayType( + inner_type=ValueType(name='float'), + extents=(4,), + ), + doc='bounding volume', + ), + StructFieldDecl( + name='bvactive', + type=ArrayType( + inner_type=ValueType(name='float'), + extents=(4,), + ), + doc='active bounding volume', + ), ), ), doc='', diff --git a/python/mujoco/structs.cc b/python/mujoco/structs.cc index 1d7cc901..0de15f76 100644 --- a/python/mujoco/structs.cc +++ b/python/mujoco/structs.cc @@ -1400,6 +1400,7 @@ PYBIND11_MODULE(_structs, m) { X(offwidth); X(offheight); X(ellipsoidinertia); + X(bvactive); #undef X py::class_ mjVisualQuality(mjVisual, "Quality"); @@ -1529,6 +1530,8 @@ PYBIND11_MODULE(_structs, m) { X(slidercrank); X(crankbroken); X(frustum); + X(bv); + X(bvactive); #undef X #define X(var) \ diff --git a/python/mujoco/structs.h b/python/mujoco/structs.h index 5fe3e1b3..5aba80f5 100644 --- a/python/mujoco/structs.h +++ b/python/mujoco/structs.h @@ -233,6 +233,8 @@ class MjWrapper : public WrapperBase { X(slidercrank); X(crankbroken); X(frustum); + X(bv); + X(bvactive); #undef X }; diff --git a/simulate/simulate.cc b/simulate/simulate.cc index d9d4f99d..f9938468 100644 --- a/simulate/simulate.cc +++ b/simulate/simulate.cc @@ -896,6 +896,7 @@ void MakeVisualizationSection(mj::Simulate* sim, const mjModel* m, int oldstate) {mjITEM_EDITNUM, "Extent", 2, &(stat->extent), "1"}, {mjITEM_EDITFLOAT, "Field of view", 2, &(vis->global.fovy), "1"}, {mjITEM_RADIO, "Inertia", 5, &(vis->global.ellipsoidinertia), "Box\nEllipsoid"}, + {mjITEM_RADIO, "BVH active", 5, &(vis->global.bvactive), "False\nTrue"}, {mjITEM_SEPARATOR, "Map", 1}, {mjITEM_EDITFLOAT, "Stiffness", 2, &(vis->map.stiffness), "1"}, {mjITEM_EDITFLOAT, "Rot stiffness", 2, &(vis->map.stiffnessrot), "1"}, diff --git a/src/engine/engine_collision_driver.c b/src/engine/engine_collision_driver.c index b8a8a7d0..e3a18c2d 100644 --- a/src/engine/engine_collision_driver.c +++ b/src/engine/engine_collision_driver.c @@ -273,7 +273,9 @@ void mj_collision(const mjModel* m, mjData* d) { mj_clearEfc(d); // reset the visualization flags - memset(d->bvh_active, 0, m->nbvh); + if (m->vis.global.bvactive) { + memset(d->bvh_active, 0, m->nbvh); + } // return if disabled if (mjDISABLED(mjDSBL_CONSTRAINT) || mjDISABLED(mjDSBL_CONTACT) @@ -634,6 +636,7 @@ void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2, mjtByte isbody2 = (bf2 < nbody); int f1 = isbody1 ? -1 : bf1 - nbody; int f2 = isbody2 ? -1 : bf2 - nbody; + int mark_active = m->vis.global.bvactive; const int bvhadr1 = isbody1 ? m->body_bvhadr[bf1] : m->flex_bvhadr[f1]; const int bvhadr2 = isbody2 ? m->body_bvhadr[bf2] : m->flex_bvhadr[f2]; const int* child1 = m->bvh_child + 2*bvhadr1; @@ -705,8 +708,10 @@ void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2, d->geom_xpos + 3*nodeid2, d->geom_xmat + 9*nodeid2, margin, NULL, NULL, &initialize)) { mj_collideGeomPair(m, d, nodeid1, nodeid2, merged, startadr, pairadr); - d->bvh_active[node1 + bvhadr1] = 1; - d->bvh_active[node2 + bvhadr2] = 1; + if (mark_active) { + d->bvh_active[node1 + bvhadr1] = 1; + d->bvh_active[node2 + bvhadr2] = 1; + } } } continue; @@ -742,8 +747,10 @@ void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2, if (m->geom_type[nodeid1] != mjGEOM_PLANE) { mj_collideGeomElem(m, d, nodeid1, f2, nodeid2); } - d->bvh_active[node1 + bvhadr1] = 1; - d->bvh_active[node2 + bvhadr2] = 1; + if (mark_active) { + d->bvh_active[node1 + bvhadr1] = 1; + d->bvh_active[node2 + bvhadr2] = 1; + } } } continue; @@ -771,8 +778,10 @@ void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2, // box filter applied in mj_collideElems, bitmask filter applied earlier if (isleaf1 && isleaf2) { mj_collideElems(m, d, f1, nodeid1, f2, nodeid2); - d->bvh_active[node1 + bvhadr1] = 1; - d->bvh_active[node2 + bvhadr2] = 1; + if (mark_active) { + d->bvh_active[node1 + bvhadr1] = 1; + d->bvh_active[node2 + bvhadr2] = 1; + } continue; } @@ -784,8 +793,10 @@ void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2, } } - d->bvh_active[node1 + bvhadr1] = 1; - d->bvh_active[node2 + bvhadr2] = 1; + if (mark_active) { + d->bvh_active[node1 + bvhadr1] = 1; + d->bvh_active[node2 + bvhadr2] = 1; + } // keep traversing the tree if (!isleaf1 && isleaf2) { diff --git a/src/engine/engine_collision_sdf.c b/src/engine/engine_collision_sdf.c index 82237438..79d97e03 100644 --- a/src/engine/engine_collision_sdf.c +++ b/src/engine/engine_collision_sdf.c @@ -496,7 +496,7 @@ static void collideBVH(const mjModel* m, mjData* d, int g, const int* faceid = m->bvh_nodeid + bvhadr; const mjtNum* bvh = m->bvh_aabb + 6*bvhadr; const int* child = m->bvh_child + 2*bvhadr; - mjtByte* visited = d->bvh_active + bvhadr; + mjtByte* bvh_active = m->vis.global.bvactive ? d->bvh_active + bvhadr : NULL; mj_markStack(d); // TODO(quaglino): Store bvh max depths to make this bound tighter. @@ -521,9 +521,6 @@ static void collideBVH(const mjModel* m, mjData* d, int g, // node1 is a leaf if (faceid[node] != -1) { - if (visited[node]) { - continue; - } if (boxIntersect(bvh+6*node, offset, rotation, m, sdf, d)) { faces[*npoints] = faceid[node]; if (++(*npoints) == MAXSDFFACE) { @@ -531,7 +528,7 @@ static void collideBVH(const mjModel* m, mjData* d, int g, mj_freeStack(d); return; } - visited[node] = 1; + if (bvh_active) bvh_active[node] = 1; } continue; } @@ -541,7 +538,7 @@ static void collideBVH(const mjModel* m, mjData* d, int g, continue; } - visited[node] = 1; + if (bvh_active) bvh_active[node] = 1; // recursive call for (int i=0; i < 2; i++) { diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index 7c44b54f..1b1975e2 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -194,6 +194,7 @@ void mj_defaultVisual(mjVisual* vis) { vis->global.offheight = 480; vis->global.realtime = 1.0; vis->global.ellipsoidinertia = 0; + vis->global.bvactive = 1; // rendering quality vis->quality.shadowsize = 4096; @@ -272,6 +273,8 @@ void mj_defaultVisual(mjVisual* vis) { setf4(vis->rgba.slidercrank, .5, .3, .8, 1.); setf4(vis->rgba.crankbroken, .9, .0, .0, 1.); setf4(vis->rgba.frustum, 1., 1., .0, .2); + setf4(vis->rgba.bv, 0., 1., .0, .5); + setf4(vis->rgba.bvactive, 1., 0., .0, .5); } diff --git a/src/engine/engine_ray.c b/src/engine/engine_ray.c index c33edd5c..b520a818 100644 --- a/src/engine/engine_ray.c +++ b/src/engine/engine_ray.c @@ -627,6 +627,7 @@ int mju_raySlab(const mjtNum aabb[6], const mjtNum xpos[3], // ray vs tree intersection mjtNum mju_rayTree(const mjModel* m, const mjData* d, int id, const mjtNum* pnt, const mjtNum* vec) { + int mark_active = m->vis.global.bvactive; const int meshid = m->geom_dataid[id]; const int bvhadr = m->mesh_bvhadr[meshid]; const int* faceid = m->bvh_nodeid + bvhadr; @@ -701,12 +702,17 @@ mjtNum mju_rayTree(const mjModel* m, const mjData* d, int id, const mjtNum* pnt, // update if (sol >= 0 && (x < 0 || sol < x)) { x = sol; + if (mark_active) { + d->bvh_active[node + bvhadr] = 1; + } } continue; } // used for rendering - d->bvh_active[node + bvhadr] = 1; + if (mark_active) { + d->bvh_active[node + bvhadr] = 1; + } // add children to the stack for (int i=0; i < 2; i++) { diff --git a/src/engine/engine_vis_visualize.c b/src/engine/engine_vis_visualize.c index d2e2de20..74bd9cd5 100644 --- a/src/engine/engine_vis_visualize.c +++ b/src/engine/engine_vis_visualize.c @@ -532,6 +532,7 @@ void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* vopt, mjvGeom* thisgeom; mjvPerturb localpert; float scl = m->stat.meansize; + int mark_active = m->vis.global.bvactive; // make default pert if missing if (!pert) { @@ -636,7 +637,6 @@ void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* vopt, category = mjCAT_DECOR; objtype = mjOBJ_UNKNOWN; if (vopt->flags[mjVIS_BODYBVH]) { - float rgba[] = {1, 0, 0, 1}; for (int i = 0; i < m->nbvhstatic; i++) { int isleaf = m->bvh_child[2*i] == -1 && m->bvh_child[2*i+1] == -1; if (scn->ngeom >= scn->maxgeom) break; @@ -671,8 +671,11 @@ void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* vopt, mju_rotVecMat(pos, center, xmat); mju_addTo3(pos, xpos); - rgba[0] = d->bvh_active[i] ? 1 : 0; - rgba[1] = d->bvh_active[i] ? 0 : 1; + // set box color + const float* rgba = m->vis.rgba.bv; + if (mark_active && d->bvh_active[i]) { + rgba = m->vis.rgba.bvactive; + } START mjv_initGeom(thisgeom, mjGEOM_LINEBOX, size, pos, xmat, rgba); @@ -685,10 +688,8 @@ void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* vopt, category = mjCAT_DECOR; objtype = mjOBJ_UNKNOWN; if (vopt->flags[mjVIS_FLEXBVH]) { - float rgba[] = {1, 0, 0, 0.1}; for (int f=0; f < m->nflex; f++) { - if (m->flex_bvhnum[f] && - vopt->flexgroup[mjMAX(0, mjMIN(mjNGROUP-1, m->flex_group[f]))]) { + if (m->flex_bvhnum[f] && vopt->flexgroup[mjMAX(0, mjMIN(mjNGROUP-1, m->flex_group[f]))]) { for (int i=m->flex_bvhadr[f]; i < m->flex_bvhadr[f]+m->flex_bvhnum[f]; i++) { int isleaf = m->bvh_child[2*i] == -1 && m->bvh_child[2*i+1] == -1; if (scn->ngeom >= scn->maxgeom) break; @@ -698,10 +699,14 @@ void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* vopt, } } - // set box data + // get box data mjtNum *aabb = d->bvh_aabb_dyn + 6*(i - m->nbvhstatic); - rgba[0] = d->bvh_active[i] ? 1 : 0; - rgba[1] = d->bvh_active[i] ? 0 : 1; + + // set box color + const float* rgba = m->vis.rgba.bv; + if (mark_active && d->bvh_active[i]) { + rgba = m->vis.rgba.bvactive; + } START mjv_initGeom(thisgeom, mjGEOM_LINEBOX, aabb+3, aabb, NULL, rgba); @@ -715,7 +720,6 @@ void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* vopt, category = mjCAT_DECOR; objtype = mjOBJ_UNKNOWN; if (vopt->flags[mjVIS_MESHBVH]) { - float rgba[] = {1, 0, 0, 1}; for (int geomid = 0; geomid < m->ngeom; geomid++) { int meshid = m->geom_dataid[geomid]; if (meshid == -1) { @@ -732,13 +736,17 @@ void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* vopt, } } - if (!d->bvh_active[i]) { - continue; + // box color + const float* rgba = m->vis.rgba.bv; + if (mark_active) { + if (d->bvh_active[i]) { + rgba = m->vis.rgba.bvactive; + } else { + // when marking active bvs, skip inactive volumes + continue; + } } - rgba[0] = d->bvh_active[i] ? 1 : 0; - rgba[1] = d->bvh_active[i] ? 0 : 1; - // get xpos, xmat, size const mjtNum* xpos = d->geom_xpos + 3 * geomid; const mjtNum* xmat = d->geom_xmat + 9 * geomid; diff --git a/src/xml/xml_native_reader.cc b/src/xml/xml_native_reader.cc index 48de27c2..aeb074fb 100644 --- a/src/xml/xml_native_reader.cc +++ b/src/xml/xml_native_reader.cc @@ -113,8 +113,8 @@ static const char* MJCF[nMJCF][mjXATTRNUM] = { {"visual", "*", "0"}, {"<"}, - {"global", "?", "10", "fovy", "ipd", "azimuth", "elevation", "linewidth", "glow", - "offwidth", "offheight", "realtime", "ellipsoidinertia"}, + {"global", "?", "11", "fovy", "ipd", "azimuth", "elevation", "linewidth", "glow", + "offwidth", "offheight", "realtime", "ellipsoidinertia", "bvactive"}, {"quality", "?", "5", "shadowsize", "offsamples", "numslices", "numstacks", "numquads"}, {"headlight", "?", "4", "ambient", "diffuse", "specular", "active"}, @@ -124,11 +124,11 @@ static const char* MJCF[nMJCF][mjXATTRNUM] = { {"scale", "?", "17", "forcewidth", "contactwidth", "contactheight", "connect", "com", "camera", "light", "selectpoint", "jointlength", "jointwidth", "actuatorlength", "actuatorwidth", "framelength", "framewidth", "constraint", "slidercrank", "frustum"}, - {"rgba", "?", "23", "fog", "haze", "force", "inertia", "joint", + {"rgba", "?", "25", "fog", "haze", "force", "inertia", "joint", "actuator", "actuatornegative", "actuatorpositive", "com", "camera", "light", "selectpoint", "connect", "contactpoint", "contactforce", "contactfriction", "contacttorque", "contactgap", "rangefinder", - "constraint", "slidercrank", "crankbroken", "frustum"}, + "constraint", "slidercrank", "crankbroken", "frustum", "bv", "bvactive"}, {">"}, {"statistic", "*", "5", "meaninertia", "meanmass", "meansize", "extent", "center"}, @@ -2723,6 +2723,10 @@ void mjXReader::Visual(XMLElement* section) { if (MapValue(elem, "ellipsoidinertia", &ellipsoidinertia, bool_map, 2)) { vis->global.ellipsoidinertia = (ellipsoidinertia==1); } + int bvactive; + if (MapValue(elem, "bvactive", &bvactive, bool_map, 2)) { + vis->global.bvactive = (bvactive==1); + } } // quality sub-element @@ -2808,6 +2812,8 @@ void mjXReader::Visual(XMLElement* section) { ReadAttr(elem, "slidercrank", 4, vis->rgba.slidercrank, text); ReadAttr(elem, "crankbroken", 4, vis->rgba.crankbroken, text); ReadAttr(elem, "frustum", 4, vis->rgba.frustum, text); + ReadAttr(elem, "bv", 4, vis->rgba.bv, text); + ReadAttr(elem, "bvactive", 4, vis->rgba.bvactive, text); } // advance to next element diff --git a/src/xml/xml_native_writer.cc b/src/xml/xml_native_writer.cc index 926e7089..88bc6fc6 100644 --- a/src/xml/xml_native_writer.cc +++ b/src/xml/xml_native_writer.cc @@ -1030,6 +1030,7 @@ void mjXWriter::Visual(XMLElement* root) { WriteAttrInt(elem, "offwidth", vis->global.offwidth, visdef.global.offwidth); WriteAttrInt(elem, "offheight", vis->global.offheight, visdef.global.offheight); WriteAttrKey(elem, "ellipsoidinertia", bool_map, 2, vis->global.ellipsoidinertia, visdef.global.ellipsoidinertia); + WriteAttrKey(elem, "bvactive", bool_map, 2, vis->global.bvactive, visdef.global.bvactive); if (!elem->FirstAttribute()) { section->DeleteChild(elem); } @@ -1122,6 +1123,8 @@ void mjXWriter::Visual(XMLElement* root) { WriteAttr(elem, "slidercrank", 4, vis->rgba.slidercrank, visdef.rgba.slidercrank); WriteAttr(elem, "crankbroken", 4, vis->rgba.crankbroken, visdef.rgba.crankbroken); WriteAttr(elem, "frustum", 4, vis->rgba.frustum, visdef.rgba.frustum); + WriteAttr(elem, "bv", 4, vis->rgba.bv, visdef.rgba.bv); + WriteAttr(elem, "bvactive", 4, vis->rgba.bvactive, visdef.rgba.bvactive); if (!elem->FirstAttribute()) { section->DeleteChild(elem); } diff --git a/test/engine/testdata/ray/stanford_bunny.xml b/test/engine/testdata/ray/stanford_bunny.xml index e8e5d8f3..916255d5 100644 --- a/test/engine/testdata/ray/stanford_bunny.xml +++ b/test/engine/testdata/ray/stanford_bunny.xml @@ -3,6 +3,13 @@ - + + + + + + + + diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index ad6bc8b9..19b0f4fc 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -4985,6 +4985,7 @@ public unsafe struct global { public int offwidth; public int offheight; public int ellipsoidinertia; + public int bvactive; } [StructLayout(LayoutKind.Sequential)] @@ -5067,6 +5068,8 @@ public unsafe struct rgba { public fixed float slidercrank[4]; public fixed float crankbroken[4]; public fixed float frustum[4]; + public fixed float bv[4]; + public fixed float bvactive[4]; } [StructLayout(LayoutKind.Sequential)] From e0864ab7f2cc3ed7f28c3a7af2ba3670c02bab40 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Fri, 19 Jan 2024 06:27:32 -0800 Subject: [PATCH 19/92] Fix formatting bugs in changelog. PiperOrigin-RevId: 599815807 Change-Id: I67d3eb2fd30c94a9d498ba605243e0b87b1d4ca8 --- doc/changelog.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 910d6048..857edacd 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -7,8 +7,8 @@ Upcoming version (not yet released) General ^^^^^^^ -1. Improved the :ref:discardvisual compiler flag, which now discards all visual-only assets. See - :ref:discardvisual for details. +1. Improved the :ref:`discardvisual` compiler flag, which now discards all visual-only assets. + See :ref:`discardvisual` for details. 2. Removed the :ref:`timer` for midphase colllision detection, it is now folded in with the narrowphase timer. This is because timing the two phases seperately required fine-grained timers inside the collision functions; these functions are so small and fast that the timer itself was incurring a measurable cost. From a23a368778050afb1cf271ef5f28f4941ca8aad9 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Fri, 19 Jan 2024 09:02:27 -0800 Subject: [PATCH 20/92] clean up rollout.cc PiperOrigin-RevId: 599849967 Change-Id: I899b28bf78bf0414b07314c24cfe6e24441eba30 --- python/mujoco/rollout.cc | 32 +++++++++++++++----------------- 1 file changed, 15 insertions(+), 17 deletions(-) diff --git a/python/mujoco/rollout.cc b/python/mujoco/rollout.cc index 05d87ece..eacdc8fe 100644 --- a/python/mujoco/rollout.cc +++ b/python/mujoco/rollout.cc @@ -12,15 +12,14 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include -#include #include #include #include -#include -#include "functions.h" +#include +#include "errors.h" #include "raw.h" +#include "structs.h" #include #include #include @@ -74,7 +73,6 @@ void _unsafe_rollout(const mjModel* m, mjData* d, int nstate, int nstep, // loop over initial states for (int s=0; s < nstate; s++) { - // set initial state if (state0) { mju_copy(d->qpos, state0 + s*nqva, nq); @@ -108,9 +106,9 @@ void _unsafe_rollout(const mjModel* m, mjData* d, int nstate, int nstep, mju_zero(d->xfrc_applied, 6*nbody); } if (!mocap) { - for (int j=0; jbody_mocapid[j]; - if (id>=0) { + if (id >= 0) { mju_copy3(d->mocap_pos+3*id, m->body_pos+3*j); mju_copy4(d->mocap_quat+4*id, m->body_quat+4*j); } @@ -174,7 +172,8 @@ mjtNum* get_array_ptr(std::optional> arg, int expected_size = nstate * nstep * dim; if (info.size != expected_size) { std::ostringstream msg; - msg << name << ".size should be " << expected_size << ", got " << info.size; + msg << name << ".size should be " << expected_size << + ", got " << info.size; throw py::value_error(msg.str()); } return static_cast(info.ptr); @@ -200,7 +199,6 @@ PYBIND11_MODULE(_rollout, pymodule) { std::optional state, std::optional sensordata ) { - const raw::MjModel* model = m.get(); raw::MjData* data = d.get(); @@ -213,7 +211,8 @@ PYBIND11_MODULE(_rollout, pymodule) { int nqva = model->nq + model->nv + model->na; mjtNum* init_state_ptr = get_array_ptr(init_state, "initial_state", nstate, 1, nqva); - mjtNum* ctrl_ptr = get_array_ptr(ctrl, "ctrl", nstate, nstep, model->nu); + mjtNum* ctrl_ptr = + get_array_ptr(ctrl, "ctrl", nstate, nstep, model->nu); mjtNum* qfrc_ptr = get_array_ptr(qfrc, "qfrc_applied", nstate, nstep, model->nv); mjtNum* xfrc_ptr = @@ -222,11 +221,11 @@ PYBIND11_MODULE(_rollout, pymodule) { get_array_ptr(mocap, "mocap", nstate, nstep, 7*model->nmocap); mjtNum* init_time_ptr = get_array_ptr(init_time, "init_time", nstate, 1, 1); - mjtNum* init_warmstart_ptr = - get_array_ptr(init_warmstart, "init_warmstart", nstate, 1, model->nv); + mjtNum* init_warmstart_ptr = get_array_ptr( + init_warmstart, "init_warmstart", nstate, 1, model->nv); mjtNum* state_ptr = get_array_ptr(state, "state", nstate, nstep, nqva); - mjtNum* sensordata_ptr = - get_array_ptr(sensordata, "sensordata", nstate, nstep, model->nsensordata); + mjtNum* sensordata_ptr = get_array_ptr(sensordata, "sensordata", nstate, + nstep, model->nsensordata); // perform rollouts { @@ -255,9 +254,8 @@ PYBIND11_MODULE(_rollout, pymodule) { py::arg("sensordata") = py::none(), py::doc(rollout_doc) ); +} } // namespace -} - -} +} // namespace mujoco::python From a1ddbdf7f8f0dd8b2c1babac6e508b94840e0a58 Mon Sep 17 00:00:00 2001 From: Kyle Bayes Date: Fri, 19 Jan 2024 09:02:37 -0800 Subject: [PATCH 21/92] Add an implementation of a cache for assets for compilation speedups. PiperOrigin-RevId: 599850003 Change-Id: I83e299e127058a5657d7b262aa98d36451b530f6 --- src/user/user_asset_cache.cc | 327 +++++++++++++++++++++++++++ src/user/user_asset_cache.h | 195 +++++++++++++++++ test/user/user_asset_cache_test.cc | 341 +++++++++++++++++++++++++++++ 3 files changed, 863 insertions(+) create mode 100644 src/user/user_asset_cache.cc create mode 100644 src/user/user_asset_cache.h create mode 100644 test/user/user_asset_cache_test.cc diff --git a/src/user/user_asset_cache.cc b/src/user/user_asset_cache.cc new file mode 100644 index 00000000..4724d748 --- /dev/null +++ b/src/user/user_asset_cache.cc @@ -0,0 +1,327 @@ +// Copyright 2024 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "user/user_asset_cache.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// adds a block of data for the asset and returns number of bytes stored +template +std::size_t mjCAsset::Add(const std::string& name, const std::vector& v) { + auto [it, inserted] = blocks_.insert({name, mjCAssetData()}); + if (!inserted) { + return 0; + } + + std::size_t n = v.size() * sizeof(T); + const uint8_t* ptr = reinterpret_cast(v.data()); + mjCAssetData& block = it->second; + + block.bytes = std::make_shared(n); + std::copy(ptr, ptr + n, block.bytes.get()); + block.nbytes = n; + nbytes_ += n; + return n; +} + +template std::size_t mjCAsset::Add(const std::string& name, + const std::vector& v); +template std::size_t mjCAsset::Add(const std::string& name, + const std::vector& v); +template std::size_t mjCAsset::Add(const std::string& name, + const std::vector& v); + + + +// fetches a block of data, sets n to size of data +template +const T* mjCAsset::Get(const std::string& name, std::size_t* n) const { + auto it = blocks_.find(name); + if (it == blocks_.end()) { + *n = 0; + return nullptr; + } + + const mjCAssetData& data = it->second; + + // TODO(kylebayes): This is probably caused by a user bug and needs an + // assertion for debugging purposes + if (data.nbytes % sizeof(T)) { + *n = 0; + return nullptr; + } + + *n = data.nbytes / sizeof(T); + return reinterpret_cast(data.bytes.get()); +} + +template +const int* mjCAsset::Get(const std::string& name, std::size_t* n) const; +template +const float* mjCAsset::Get(const std::string& name, std::size_t* n) const; +template +const double* mjCAsset::Get(const std::string& name, std::size_t* n) const; + + + +// replaces blocks data in asset +void mjCAsset::ReplaceBlocks( + const std::unordered_map& blocks, + std::size_t nbytes) { + blocks_ = blocks; + nbytes_ = nbytes; +} + + + +// makes a copy for user (strip unnecessary items) +mjCAsset mjCAsset::Copy(const mjCAsset& other) { + mjCAsset asset; + asset.id_ = other.Id(); + asset.timestamp_ = other.Timestamp(); + asset.blocks_ = other.blocks_; + asset.nbytes_ = other.nbytes_; + return asset; +} + + + +// sets the total maximum size of the cache in bytes +// low-priority cached assets will be dropped to make the new memory +// requirement +void mjCCache::SetMaxSize(std::size_t size) { + std::lock_guard lock(mutex_); + max_size_ = size; + Trim(); +} + + + +// returns the corresponding timestamp, if the given asset is stored in the cache +const std::string* mjCCache::HasAsset(const std::string& id) { + std::lock_guard lock(mutex_); + auto it = lookup_.find(id); + if (it == lookup_.end()) { + return nullptr; + } + + return &(it->second.Timestamp()); +} + + + +// inserts an asset into the cache, if asset is already in the cache, its data +// is updated only if the timestamps disagree +bool mjCCache::Insert(const mjCAsset& asset) { + std::lock_guard lock(mutex_); + const std::string& id = asset.Id(); + if (asset.References().size() != 1) { + return false; + } + const std::string& filename = *(asset.References().begin()); + auto [it, inserted] = lookup_.insert({id, asset}); + + if (!inserted) { + mjCAsset* asset_ptr = &(it->second); + if (size_ - asset_ptr->BytesCount() + asset.BytesCount() > max_size_) { + return false; + } + models_[filename].insert(asset_ptr); // add it for the model + asset_ptr->AddReference(filename); + if (it->second.Timestamp() == asset.Timestamp()) { + return true; + } + asset_ptr->SetTimestamp(asset.Timestamp()); + size_ = size_ - asset_ptr->BytesCount() + asset.BytesCount(); + asset_ptr->ReplaceBlocks(asset.Blocks(), asset.BytesCount()); + return true; + } else if (size_ + asset.BytesCount() > max_size_) { + return false; + } + + // new asset + mjCAsset* asset_ptr = &(it->second); + asset_ptr->SetInsertNum(insert_num_++); + entries_.insert(asset_ptr); + models_[filename].insert(asset_ptr); + size_ += asset.BytesCount(); + return true; +} + + + +bool mjCCache::Insert(mjCAsset&& asset) { + std::lock_guard lock(mutex_); + const std::string& id = asset.Id(); + if (asset.References().size() != 1) { + return false; + } + const std::string& filename = *(asset.References().begin()); + std::size_t nbytes = asset.BytesCount(); + auto [it, inserted] = lookup_.try_emplace(id, std::move(asset)); + + if (!inserted) { + mjCAsset* asset_ptr = &(it->second); + if (size_ - asset_ptr->BytesCount() + nbytes > max_size_) { + return false; + } + models_[filename].insert(asset_ptr); // add it for the model + asset_ptr->AddReference(std::move(filename)); + if (it->second.Timestamp() == asset.Timestamp()) { + return true; + } + // move data and timestamp over + asset_ptr->SetTimestamp(std::move(asset.timestamp_)); + size_ = size_ - asset_ptr->BytesCount() + nbytes; + asset_ptr->ReplaceBlocks(std::move(asset.blocks_), asset.nbytes_); + return true; + } else if (size_ + nbytes > max_size_) { + return false; + } + + // new asset + mjCAsset* asset_ptr = &(it->second); + asset_ptr->SetInsertNum(insert_num_++); + entries_.insert(asset_ptr); + models_[filename].insert(asset_ptr); + size_ += nbytes; + return true; +} + + + +// returns the asset with the given id, if it exists in the cache +std::optional mjCCache::Get(const std::string& id) { + std::lock_guard lock(mutex_); + auto it = lookup_.find(id); + if (it == lookup_.end()) { + return std::nullopt; + } + + mjCAsset* asset = &(it->second); + asset->IncrementAccess(); + + // update priority queue + entries_.erase(asset); + entries_.insert(asset); + return asset->Copy(*asset); +} + + + +// removes model from the cache along with assets referencing only this model +void mjCCache::RemoveModel(const std::string& filename) { + std::lock_guard lock(mutex_); + for (mjCAsset* asset : models_[filename]) { + asset->RemoveReference(filename); + if (!asset->HasReferences()) { + Delete(asset, filename); + } + } + models_.erase(filename); +} + + + +// Wipes out all internal data for the given model +void mjCCache::Reset(const std::string& filename) { + std::lock_guard lock(mutex_); + for (auto asset : models_[filename]) { + Delete(asset, filename); + } + models_.erase(filename); +} + + + +// Wipes out all internal data +void mjCCache::Reset() { + std::lock_guard lock(mutex_); + entries_.clear(); + lookup_.clear(); + models_.clear(); + size_ = 0; + insert_num_ = 0; +} + + + +std::size_t mjCCache::MaxSize() const { + std::lock_guard lock(mutex_); + return max_size_; +} + + + +std::size_t mjCCache::Size() const { + std::lock_guard lock(mutex_); + return size_; +} + + + +// Deletes a single asset +void mjCCache::DeleteAsset(const std::string& id) { + std::lock_guard lock(mutex_); + auto it = lookup_.find(id); + if (it != lookup_.end()) { + Delete(&(it->second)); + } +} + + + +// Deletes a single asset (internal) +void mjCCache::Delete(mjCAsset* asset) { + size_ -= asset->BytesCount(); + entries_.erase(asset); + for (auto& reference : asset->References()) { + models_[reference].erase(asset); + } + lookup_.erase(asset->Id()); +} + + + +// Deletes a single asset (internal) +void mjCCache::Delete(mjCAsset* asset, const std::string& skip) { + size_ -= asset->BytesCount(); + entries_.erase(asset); + + for (auto& reference : asset->References()) { + if (reference != skip) { + models_[reference].erase(asset); + } + } + lookup_.erase(asset->Id()); +} + + + +// trims out data to meet memory requirements +void mjCCache::Trim() { + while (size_ > max_size_) { + Delete(*entries_.begin()); + } +} diff --git a/src/user/user_asset_cache.h b/src/user/user_asset_cache.h new file mode 100644 index 00000000..e8f0bc7c --- /dev/null +++ b/src/user/user_asset_cache.h @@ -0,0 +1,195 @@ +// Copyright 2024 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MUJOCO_SRC_USER_ASSET_CACHE_H_ +#define MUJOCO_SRC_USER_ASSET_CACHE_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// data associated with an asset +struct mjCAssetData { + std::shared_ptr bytes; // raw serialized bytes of cached data + std::size_t nbytes; // number of bytes stored +}; + +// A class container for a thread-safe asset cache +// +// Each mjCAsset is used to store raw and/or processed data loaded from a +// resource and is defined by a unique ID (usually the full filename of the +// asset). The asset's data can be segregated into blocks for ease of use of +// mix and matching different types of data. Each block is given a unique name +// within the asset for readability. For example, mjCAsset for a mesh may +// include all vertex positions, edges, and the computed volume. +class mjCAsset { + friend class mjCCache; + public: + mjCAsset(std::string filename, std::string id, std::string timestamp) + : id_(std::move(id)), timestamp_(std::move(timestamp)) { + AddReference(filename); + } + + // move and copy constructors + mjCAsset(mjCAsset&& other) = default; + mjCAsset& operator=(mjCAsset&& other) = default; + mjCAsset(const mjCAsset& other) = default; + mjCAsset& operator=(const mjCAsset& other) = default; + + // adds a block of data for the asset and returns number of bytes stored + // loading data into an asset should happen single thread + template std::size_t Add(const std::string& name, + const std::vector& v); + + // fetches a block of data, sets n to size of data + // TODO(kylebayes): The C++ span utility doesn't seem to be supported by + // Google C++ coding standards. For now, we fallback to an C style API. + template + const T* Get(const std::string& name, std::size_t* n) const; + + private: + mjCAsset() = default; + + // helpers for managing models referencing this asset + void AddReference(std::string xml_file) { references_.insert(xml_file); } + void RemoveReference(const std::string& xml_file) { + references_.erase(xml_file); + } + bool HasReferences() const { return !references_.empty(); } + + // replaces data blocks in asset + void ReplaceBlocks( + const std::unordered_map& blocks, + std::size_t nbytes); + + void IncrementAccess() { access_count_++; } + + // makes a copy for user (strip unnecessary references) + static mjCAsset Copy(const mjCAsset& other); + + // setters + void SetInsertNum(std::size_t num) { insert_num_ = num; } + void SetTimestamp(std::string timestamp) { timestamp_ = timestamp; } + + // accessors + const std::string& Id() const { return id_; } + const std::string& Timestamp() const { return timestamp_; } + std::size_t InsertNum() const { return insert_num_; } + std::size_t AccessCount() const { return access_count_; } + std::size_t BytesCount() const { return nbytes_; } + const std::unordered_map& Blocks() const { + return blocks_; + } + const std::set& References() const { return references_; } + + std::string id_; // unique id associated with asset + std::string timestamp_; // opaque timestamp of asset + std::size_t insert_num_; // number when asset was inserted + std::size_t access_count_ = 0; // incremented when getting 0th block + std::size_t nbytes_ = 0; // how many bytes taken up by the asset + + // the actually data of the asset + std::unordered_map blocks_; + + // list of models referencing this asset + std::set references_; +}; + +// the class container for a thread-safe asset cache +class mjCCache { + public: + explicit mjCCache(std::size_t size) : max_size_(size) {} + + // move only + mjCCache(mjCCache&& other) = default; + mjCCache& operator=(mjCCache&& other) = default; + mjCCache(const mjCCache& other) = delete; + mjCCache& operator=(const mjCCache& other) = delete; + + // sets the total maximum size of the cache in bytes + // low-priority cached assets will be dropped to make the new memory + // requirement + void SetMaxSize(std::size_t size); + + // returns the corresponding timestamp, if the given asset is stored in + // the cache + const std::string* HasAsset(const std::string& id); + + // inserts an asset into the cache, if asset is already in the cache, its data + // is updated only if the timestamps disagree + bool Insert(const mjCAsset& asset); + bool Insert(mjCAsset&& asset); + + // returns the asset with the given id, if it exists in the cache + std::optional Get(const std::string& id); + + // deletes the asset from the cache with the given id + void DeleteAsset(const std::string& id); + + // removes model from the cache, assets only referenced by the model will be + // deleted + void RemoveModel(const std::string& filename); + + // Wipes out all assets from the cache for the given model + void Reset(const std::string& filename); + + // Wipes out all internal data + void Reset(); + + // accessors + std::size_t MaxSize() const; + std::size_t Size() const; + + private: + void Delete(mjCAsset* asset); + void Delete(mjCAsset* asset, const std::string& skip); + void Trim(); + + // TODO(kylebayes): We should consider a shared mutex like in + // engine/engine_plugin.cc as some of these methods don't need to be fully + // locked. + mutable std::mutex mutex_; + std::size_t insert_num_ = 0; // a running counter of assets being inserted + std::size_t size_ = 0; // current size of the cache in bytes + std::size_t max_size_ = 0; // max size of the cache in bytes + + // compare function for the priority queue + static constexpr auto compare_ = [](const mjCAsset* e1, + const mjCAsset* e2) { + if (e1->AccessCount() != e2->AccessCount()) { + return e1->AccessCount() < e2->AccessCount(); + } + return e1->InsertNum() < e2->InsertNum(); + }; + + // internal constant look up table for assets + std::unordered_map lookup_; + + // internal priority queue for the cache + std::set entries_; + + // models using the cache along with the assets they reference + std::unordered_map> models_; +}; + +#endif // MUJOCO_SRC_USER_ASSET_CACHE_H_ diff --git a/test/user/user_asset_cache_test.cc b/test/user/user_asset_cache_test.cc new file mode 100644 index 00000000..577380de --- /dev/null +++ b/test/user/user_asset_cache_test.cc @@ -0,0 +1,341 @@ +// Copyright 2024 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Tests for user/user_asset_cache.cc + +#include +#include +#include +#include + +#include +#include +#include "test/fixture.h" +#include "src/user/user_asset_cache.h" + +namespace mujoco { + +using ::testing::ElementsAreArray; +using ::testing::IsNull; +using ::testing::NotNull; +using ::testing::StrEq; + +using AssetCacheTest = MujocoTest; + +namespace { + +constexpr int kMaxSize = 100; // in bytes + +TEST(AssetCacheTest, SizeTest) { + mjCCache cache(kMaxSize); + EXPECT_EQ(cache.Size(), 0); +} + +TEST(AssetCacheTest, HasAssetSuccessTest) { + mjCCache cache(kMaxSize); + + mjCAsset asset("file.xml", "foo.obj", "now"); + cache.Insert(asset); + + EXPECT_THAT(*(cache.HasAsset("foo.obj")), StrEq("now")); +} + +TEST(AssetCacheTest, HasAssetFailureTest) { + mjCCache cache(kMaxSize); + + mjCAsset asset("file.xml", "foo.obj", "now"); + cache.Insert(asset); + + EXPECT_THAT(cache.HasAsset("file2.xml"), nullptr); +} + +TEST(AssetCacheTest, AddSuccessTest) { + mjCCache cache(kMaxSize); + std::vector v1 = {1, 2, 3}; + std::vector v2 = {1.0, 2.0, 3.0}; + + mjCAsset asset("file.xml", "foo.obj", "now"); + std::size_t nbytes1 = asset.Add("v1", v1); + std::size_t nbytes2 = asset.Add("v2", v2); + cache.Insert(asset); + + ASSERT_EQ(nbytes1, 12); + ASSERT_EQ(nbytes2, 24); + ASSERT_EQ(cache.Size(), 36); +} + +TEST(AssetCacheTest, AddFailureTest) { + mjCCache cache(kMaxSize); + std::vector v1 = {1, 2, 3}; + std::vector v2 = {1.0, 2.0, 3.0}; + + mjCAsset asset("file.xml", "foo.obj", "now"); + asset.Add("v1", v1); + std::size_t nbytes = asset.Add("v1", v2); + cache.Insert(asset); + + ASSERT_EQ(nbytes, 0); + ASSERT_EQ(cache.Size(), 12); +} + +TEST(AssetCacheTest, InsertReplaceTest) { + mjCCache cache(kMaxSize); + std::vector v1 = {1, 2, 3}; + std::vector v2 = {1.0, 2.0, 3.0}; + + mjCAsset asset("file.xml", "foo.obj", "now"); + asset.Add("v", v1); + cache.Insert(asset); + + mjCAsset asset2("file.xml", "foo.obj", "nower"); + asset2.Add("v", v2); + bool inserted = cache.Insert(asset2); + EXPECT_TRUE(inserted); + + mjCAsset asset3 = *(cache.Get("foo.obj")); + std::size_t n = 0; + const double* ptr = asset3.Get("v", &n); + std::vector v3 = std::vector(ptr, ptr + n); + EXPECT_THAT(v3, ElementsAreArray(v2)); + + ASSERT_EQ(cache.Size(), 24); +} + +TEST(AssetCacheTest, MoveInsertNewTest) { + mjCCache cache(kMaxSize); + std::vector v1 = {1, 2, 3}; + std::vector v2 = {1.0, 2.0, 3.0}; + + mjCAsset asset("file.xml", "foo.obj", "now"); + std::size_t nbytes1 = asset.Add("v1", v1); + std::size_t nbytes2 = asset.Add("v2", v2); + cache.Insert(std::move(asset)); + + ASSERT_EQ(nbytes1, 12); + ASSERT_EQ(nbytes2, 24); + ASSERT_EQ(cache.Size(), 36); +} + +TEST(AssetCacheTest, MoveInsertReplaceTest) { + mjCCache cache(kMaxSize); + std::vector v1 = {1, 2, 3}; + std::vector v2 = {1.0, 2.0, 3.0}; + + mjCAsset asset("file.xml", "foo.obj", "now"); + asset.Add("v", v1); + cache.Insert(std::move(asset)); + + mjCAsset asset2("file.xml", "foo.obj", "nower"); + asset2.Add("v", v2); + bool inserted = cache.Insert(std::move(asset2)); + EXPECT_TRUE(inserted); + + mjCAsset asset3 = *(cache.Get("foo.obj")); + std::size_t n = 0; + const double* ptr = asset3.Get("v", &n); + std::vector v3 = std::vector(ptr, ptr + n); + EXPECT_THAT(v3, ElementsAreArray(v2)); + + ASSERT_EQ(cache.Size(), 24); +} + +TEST(AssetCacheTest, GetSuccessTest) { + mjCCache cache(kMaxSize); + std::vector v1 = {1, 2, 3}; + std::vector v2 = {1.0, 2.0, 3.0}; + mjCAsset asset("file.xml", "foo.obj", "now"); + asset.Add("v1", v1); + asset.Add("v2", v2); + cache.Insert(asset); + + std::size_t n = 0; + + mjCAsset asset2 = *(cache.Get("foo.obj")); + + const int* ptr1 = asset2.Get("v1", &n); + EXPECT_EQ(n, 3); + std::vector v3 = std::vector(ptr1, ptr1 + n); + EXPECT_THAT(v3, ElementsAreArray(v1)); + + const double* ptr2 = asset2.Get("v2", &n); + EXPECT_EQ(n, 3); + std::vector v4 = std::vector(ptr2, ptr2 + n); + EXPECT_THAT(v4, ElementsAreArray(v3)); +} + +TEST(AssetCacheTest, GetFailueTest) { + mjCCache cache(kMaxSize); + std::vector v = {1, 2, 3}; + mjCAsset asset("file.xml", "foo.obj", "now"); + asset.Add("v", v); + cache.Insert(asset); + + mjCAsset asset2 = *(cache.Get("foo.obj")); + EXPECT_EQ(cache.Get("bar.obj").has_value(), false); + + std::size_t n = 0; + const int* ptr = asset.Get("v2", &n); + EXPECT_THAT(ptr, IsNull()); +} + +// Trim cache based off of access count +TEST(AssetCacheTest, LimitTest1) { + mjCCache cache(kMaxSize); + EXPECT_THAT(cache.MaxSize(), kMaxSize); + std::vector v = {1, 2, 3}; + + mjCAsset asset1 = mjCAsset("file.xml", "foo.obj", "now"); + mjCAsset asset2 = mjCAsset("file.xml", "bar.obj", "now"); + asset1.Add("foo.obj", v); + asset2.Add("bar.obj", v); + cache.Insert(asset1); + cache.Insert(asset2); + + // access asset foo twice, bar one + cache.Get("foo.obj"); + cache.Get("foo.obj"); + cache.Get("bar.obj"); + + // make max size so cache can hold only one asset + cache.SetMaxSize(12); + + // foo should still be in cache + EXPECT_THAT(cache.HasAsset("foo.obj"), NotNull()); + + // bar was accessed less, so is removed + EXPECT_THAT(cache.HasAsset("bar.obj"), IsNull()); +} + +// Trim cache based off of insert order +TEST(AssetCacheTest, LimitTest2) { + mjCCache cache(kMaxSize); + std::vector v = {1, 2, 3}; + mjCAsset asset1("file.xml", "foo.obj", "now"); + mjCAsset asset2("file.xml", "bar.obj", "now"); + asset1.Add("v", v); + asset2.Add("v", v); + + cache.Insert(asset1); + cache.Insert(asset2); + + // get each asset once + mjCAsset asset3 = *(cache.Get("foo.obj")); + mjCAsset asset4 = *(cache.Get("bar.obj")); + + // make max size so cache can hold only one asset + cache.SetMaxSize(12); + + // foo should be gone because it's older + EXPECT_THAT(cache.HasAsset("foo.obj"), IsNull()); + + // bar should still be in cache + EXPECT_THAT(cache.HasAsset("bar.obj"), NotNull()); +} + +TEST(AssetCacheTest, ResetAllTest) { + mjCCache cache(kMaxSize); + mjCAsset asset1("file1.xml", "foo.obj", "now"); + mjCAsset asset2("file2.xml", "bar.obj", "now"); + cache.Insert(asset1); + cache.Insert(asset2); + + EXPECT_THAT(cache.HasAsset("foo.obj"), NotNull()); + EXPECT_THAT(cache.HasAsset("bar.obj"), NotNull()); + + cache.Reset(); + + EXPECT_THAT(cache.HasAsset("foo.obj"), IsNull()); + EXPECT_THAT(cache.HasAsset("bar.obj"), IsNull()); +} + +TEST(AssetCacheTest, ResetModelTest1) { + mjCCache cache(kMaxSize); + mjCAsset asset("file1.xml", "foo.obj", "now"); + mjCAsset asset2("file1.xml", "bar.obj", "now"); + mjCAsset asset3("file2.xml", "bar.obj", "now"); + cache.Insert(asset); + cache.Insert(asset2); + cache.Insert(asset3); + + cache.Reset("file2.xml"); + + EXPECT_THAT(cache.HasAsset("foo.obj"), NotNull()); + EXPECT_THAT(cache.HasAsset("bar.obj"), IsNull()); +} + +TEST(AssetCacheTest, ResetModelTest2) { + mjCCache cache(kMaxSize); + mjCAsset asset("file1.xml", "foo.obj", "now"); + mjCAsset asset2("file2.xml", "foo.obj", "now"); + mjCAsset asset3("file2.xml", "bar.obj", "now"); + cache.Insert(asset); + cache.Insert(asset2); + cache.Insert(asset3); + + cache.Reset("file2.xml"); + + EXPECT_THAT(cache.HasAsset("foo.obj"), IsNull()); + EXPECT_THAT(cache.HasAsset("bar.obj"), IsNull()); +} + +TEST(AssetCacheTest, RemoveModelTest1) { + mjCCache cache(kMaxSize); + mjCAsset asset("file1.xml", "foo.obj", "now"); + mjCAsset asset2("file1.xml", "bar.obj", "now"); + mjCAsset asset3("file2.xml", "bar.obj", "now"); + cache.Insert(asset); + cache.Insert(asset2); + cache.Insert(asset3); + + cache.RemoveModel("file2.xml"); + + EXPECT_THAT(cache.HasAsset("foo.obj"), NotNull()); + EXPECT_THAT(cache.HasAsset("bar.obj"), NotNull()); +} + +TEST(AssetCacheTest, RemoveModelTest2) { + mjCCache cache(kMaxSize); + mjCAsset asset("file1.xml", "foo.obj", "now"); + mjCAsset asset2("file2.xml", "bar.obj", "now"); + cache.Insert(asset); + cache.Insert(asset2); + + cache.Reset("file2.xml"); + + EXPECT_THAT(cache.HasAsset("foo.obj"), NotNull()); + EXPECT_THAT(cache.HasAsset("bar.obj"), IsNull()); +} + +TEST(AssetCacheTest, DeleteAssetSuccessTest) { + mjCCache cache(kMaxSize); + mjCAsset asset("file1.xml", "foo.obj", "now"); + cache.Insert(asset); + + cache.DeleteAsset("foo.obj"); + + EXPECT_THAT(cache.HasAsset("foo.obj"), IsNull()); +} + +TEST(AssetCacheTest, DeleteAssetFailureTest) { + mjCCache cache(kMaxSize); + mjCAsset asset("file.xml", "foo.obj", "now"); + cache.Insert(asset); + + cache.DeleteAsset("bar.obj"); + + EXPECT_THAT(cache.HasAsset("foo.obj"), NotNull()); +} + +} // namespace +} // namespace mujoco From 5cbaa233882b187c05bb004422a6b653db6f2c0b Mon Sep 17 00:00:00 2001 From: Erik Frey Date: Fri, 19 Jan 2024 11:14:06 -0800 Subject: [PATCH 22/92] Remove TEST_FILES from MJX test_util.py. PiperOrigin-RevId: 599886327 Change-Id: I4d166e09db22fffc959dbe34fa3469bd4a1b9f84 --- mjx/mujoco/mjx/_src/device_test.py | 6 ++--- mjx/mujoco/mjx/_src/support_test.py | 2 +- mjx/mujoco/mjx/_src/test_util.py | 9 +------ mjx/mujoco/mjx/_src/test_util_test.py | 38 --------------------------- 4 files changed, 5 insertions(+), 50 deletions(-) delete mode 100644 mjx/mujoco/mjx/_src/test_util_test.py diff --git a/mjx/mujoco/mjx/_src/device_test.py b/mjx/mujoco/mjx/_src/device_test.py index b4b9a0cd..71137c5d 100644 --- a/mjx/mujoco/mjx/_src/device_test.py +++ b/mjx/mujoco/mjx/_src/device_test.py @@ -59,7 +59,7 @@ def _assert_eq(testcase, a, b, attr=None, name=None): class DeviceTest(parameterized.TestCase): - @parameterized.parameters(test_util.TEST_FILES) + @parameterized.parameters('constraints.xml', 'pendula.xml') def testdevice_put(self, fname): """Test putting MjData and MjModel on device.""" m = test_util.load_test_file(fname) @@ -71,7 +71,7 @@ class DeviceTest(parameterized.TestCase): _assert_eq(self, mjx.device_put(d), d) _assert_eq(self, mjx.device_put(m), m) - @parameterized.parameters(test_util.TEST_FILES) + @parameterized.parameters('constraints.xml', 'pendula.xml') def testdevice_get(self, fname): """Test getting MjData from a device.""" m = test_util.load_test_file(fname) @@ -81,7 +81,7 @@ class DeviceTest(parameterized.TestCase): device.device_get_into(d, dx) _assert_eq(self, dx, d) - @parameterized.parameters(set(test_util.TEST_FILES) - {'convex.xml'}) + @parameterized.parameters('constraints.xml', 'pendula.xml') def testdevice_get_batched(self, fname): """Test getting MjData from a device.""" m = test_util.load_test_file(fname) diff --git a/mjx/mujoco/mjx/_src/support_test.py b/mjx/mujoco/mjx/_src/support_test.py index fb3a5389..f6a79488 100644 --- a/mjx/mujoco/mjx/_src/support_test.py +++ b/mjx/mujoco/mjx/_src/support_test.py @@ -27,7 +27,7 @@ import numpy as np class SupportTest(parameterized.TestCase): - @parameterized.parameters(set(test_util.TEST_FILES) - {'convex.xml'}) + @parameterized.parameters('constraints.xml', 'pendula.xml') def test_jac(self, fname): np.random.seed(0) diff --git a/mjx/mujoco/mjx/_src/test_util.py b/mjx/mujoco/mjx/_src/test_util.py index d350c11a..b360fb63 100644 --- a/mjx/mujoco/mjx/_src/test_util.py +++ b/mjx/mujoco/mjx/_src/test_util.py @@ -15,20 +15,13 @@ """Utilities for testing.""" import sys -from typing import Dict, List, Tuple +from typing import Dict, Tuple from xml.etree import ElementTree as ET from etils import epath import mujoco import numpy as np -TEST_FILES: List[str] = [ - 'constraints.xml', - 'convex.xml', - 'pendula.xml', - 'ray.xml', -] - _ACTUATOR_TYPES = ['motor', 'velocity', 'position', 'general', 'intvelocity'] _DYN_TYPES = ['none', 'integrator', 'filter', 'filterexact'] _DYN_PRMS = ['0.189', '2.1'] diff --git a/mjx/mujoco/mjx/_src/test_util_test.py b/mjx/mujoco/mjx/_src/test_util_test.py deleted file mode 100644 index e8d32dab..00000000 --- a/mjx/mujoco/mjx/_src/test_util_test.py +++ /dev/null @@ -1,38 +0,0 @@ -# Copyright 2023 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== -"""Tests for the test_util.""" - -from absl.testing import absltest -from etils import epath -from mujoco.mjx._src import test_util - - -class TestUtilTest(absltest.TestCase): - - def test_files_in_test_data_match(self): - directory = epath.resource_path('mujoco.mjx') / 'test_data' - files = set([f.name for f in directory.glob('*.xml')]) - self.assertSetEqual( - files, - set(test_util.TEST_FILES), - msg=( - '`_test_util.TEST_FILES` must match the files in the ' - 'test_data/*.xml directory' - ), - ) - - -if __name__ == '__main__': - absltest.main() From 042e3c9c4ce12214f07a2d04ea526810f22bfe0c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?B=C3=A1lint=20Hodossy?= Date: Sat, 20 Jan 2024 19:06:15 +0000 Subject: [PATCH 23/92] Reuse indices in mesh parsing, remove trailing nulls in mesh writing --- unity/Editor/Importer/StlMeshParser.cs | 73 +++++++++++++++++--------- 1 file changed, 49 insertions(+), 24 deletions(-) diff --git a/unity/Editor/Importer/StlMeshParser.cs b/unity/Editor/Importer/StlMeshParser.cs index 207218dc..f1dfaa5e 100644 --- a/unity/Editor/Importer/StlMeshParser.cs +++ b/unity/Editor/Importer/StlMeshParser.cs @@ -12,11 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. -using System; using System.Collections.Generic; using System.IO; using System.Linq; using UnityEngine; +using UnityEngine.Rendering; namespace Mujoco { @@ -27,6 +27,17 @@ public static class BinaryReaderExtensions { var z = reader.ReadSingle(); return new Vector3(x, y, z); } + + public static int GetOrCreateVertexIndex(Dictionary vertexIndexMap, List vertices, Listnormals, Vector3 vertex, Vector3 normal) { + if (vertexIndexMap.TryGetValue(vertex, out int existingIndex)) { + return existingIndex; + } + int newIndex = vertexIndexMap.Count; + vertexIndexMap.Add(vertex, newIndex); + vertices.Add(vertex); + normals.Add(normal); + return newIndex; + } } public static class BinaryWriterExtensions { @@ -50,7 +61,7 @@ public class StlMeshParser { // The binary STL format is described here: https://en.wikipedia.org/wiki/STL_(file_format) public static Mesh ParseBinary(byte[] stlFileContents, Vector3 scale) { var fileTypeId = System.Text.Encoding.UTF8.GetString( - stlFileContents.Take(_asciiFileTypeId.Length).ToArray()); + stlFileContents.Take(_asciiFileTypeId.Length).ToArray()); if (fileTypeId == _asciiFileTypeId) { throw new IOException("Ascii STL file format is not supported."); } @@ -59,27 +70,39 @@ public class StlMeshParser { using (var reader = new BinaryReader(stream)) { reader.ReadBytes(_headerLength); var numTriangles = reader.ReadUInt32(); - var numVertices = numTriangles * _verticesPerTriangle; - if (numVertices > _unityLimitNumVerticesPerMesh) { - throw new IndexOutOfRangeException( - "The mesh exceeds the number of vertices per mesh allowed by Unity. " + - $"({numVertices} > {_unityLimitNumVerticesPerMesh})"); - } - var triangleIndices = new List(capacity: (int)numVertices); - var vertices = new List(capacity: (int)numVertices); - var normals = new List(capacity: (int)numVertices); - for (var i = 0; i < numVertices; i += _verticesPerTriangle) { + var maxNumVertices = numTriangles * _verticesPerTriangle; + + Dictionary vertexIndexMap = new Dictionary(); + var triangleIndices = new int[(int)numTriangles * _verticesPerTriangle]; + var vertices = new List(capacity: (int)maxNumVertices); + var normals = new List(capacity: (int)maxNumVertices); + for (var i = 0; i < numTriangles; i++) { var triangleNormal = ToXZY(reader.ReadVector3()); - normals.AddRange(new[] { triangleNormal, triangleNormal, triangleNormal }); - vertices.AddRange(new[] { - ToXZY(reader.ReadVector3()), - ToXZY(reader.ReadVector3()), - ToXZY(reader.ReadVector3()) }); - triangleIndices.AddRange(new[] {i, i + 2, i + 1}); - reader.ReadInt16(); // Read the unused attribute indices field. + var verts = new[] + { + ToXZY(reader.ReadVector3()), + ToXZY(reader.ReadVector3()), + ToXZY(reader.ReadVector3()) + }; + var indices = new[] { verts[0], verts[2], verts[1] }.Select(v => + BinaryReaderExtensions.GetOrCreateVertexIndex(vertexIndexMap, + vertices, + normals, + v, + triangleNormal)).ToArray(); + for (int j = 0; j < 3; j++) { + triangleIndices[i * 3 + j] = indices[j]; + } + reader.ReadInt16(); // Read the unused attribute indices field. } var mesh = new Mesh(); + var numVertices = vertexIndexMap.Count; + + if (numVertices > _unityLimitNumVerticesPerMesh) { + mesh.indexFormat = IndexFormat.UInt32; + } + mesh.vertices = vertices.ToArray(); mesh.normals = normals.ToArray(); mesh.triangles = triangleIndices.ToArray(); @@ -88,6 +111,7 @@ public class StlMeshParser { mesh.RecalculateNormals(); mesh.RecalculateTangents(); mesh.RecalculateBounds(); + return mesh; } } @@ -118,15 +142,16 @@ public class StlMeshParser { var i2 = triangles[i + 1]; var i3 = triangles[i + 2]; var faceNormal = (normals[i1] + normals[i2] + normals[i3]).normalized; - writer.Write(ToXZY(faceNormal)); + writer.Write(faceNormal); - writer.Write(ToXZY(vertices[i1])); - writer.Write(ToXZY(vertices[i3])); - writer.Write(ToXZY(vertices[i2])); + writer.Write(vertices[i1]); + writer.Write(vertices[i2]); + writer.Write(vertices[i3]); writer.Write((short)0); } - return stream.GetBuffer(); + + return stream.ToArray(); } } } From aceb52bd0997b7970b6989e85b6fe2ab4f3aa257 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Mon, 22 Jan 2024 07:01:01 -0800 Subject: [PATCH 24/92] Improvements to `mujoco.rollout`: - `mjSTATE_FULLPHYSICS` as state spec, enabling divergence detection by inspecting time. - User-defined control spec. - Stop squeezing: outputs always have dim=3. PiperOrigin-RevId: 600445256 Change-Id: I4466e88929cb7081e1c94968a5cfe10485bb7475 --- doc/changelog.rst | 29 +- doc/python.rst | 30 +- python/mujoco/rollout.cc | 243 ++++++++--------- python/mujoco/rollout.py | 189 +++++++------ python/mujoco/rollout_test.py | 500 ++++++++++++++++++++-------------- 5 files changed, 564 insertions(+), 427 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 857edacd..cb27d8b4 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -27,12 +27,21 @@ MJX 6. Updated MJX colab tutorial with more stable quadruped environment. 7. Added ``mjx.ray`` which mirrors :ref:`mj_ray` for planes, spheres, capsules, boxes, and meshes. +Python bindings +^^^^^^^^^^^^^^^ +8. Improved the implmentation of the :ref:`rollout` module. Note the changes below are breaking, dependent + code will require modification. + + - Uses :ref:`mjSTATE_FULLPHYSICS` as state spec, enabling divergence detection by inspecting time. + - Allows user-defined control spec for any combination of :ref:`user input` fields as controls. + - Outputs are no longer squeezed and always have dim=3. + Bug fixes ^^^^^^^^^ -8. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes +9. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes :github:issue:`1270`. -9. Fixed a bug in the :ref:`muscle model` that led to non-zero values outside the lower - bound of the length range. Fixes :github:issue:`1342`. +10. Fixed a bug in the :ref:`muscle model` that led to non-zero values outside the lower + bound of the length range. Fixes :github:issue:`1342`. Version 3.1.1 (December 18, 2023) @@ -40,7 +49,8 @@ Version 3.1.1 (December 18, 2023) Bug fixes ^^^^^^^^^ -1. Fixed a bug (introduced in 3.1.0) where box-box collisions produced no contacts if one box was deeply embedded in the other. +1. Fixed a bug (introduced in 3.1.0) where box-box collisions produced no contacts if one box was deeply embedded in the + other. 2. Fixed a bug in :ref:`simulate` where the "LOADING..." message was not showing correctly. 3. Fixed a crash in the Python :ref:`passive viewer`, when used with models containing Flex objects. 4. Fixed a bug in MJX where ``site_xmat`` was ignored in ``get_data`` and ``put_data`` @@ -53,8 +63,8 @@ Version 3.1.0 (December 12, 2023) General ^^^^^^^ 1. Improved convergence of Signed Distance Function (SDF) collisions by using line search and a new objective function - for the optimization. This allows to decrease the number of initial points needed for finding the contacts and is more - robust for very small or large geom sizes. + for the optimization. This allows to decrease the number of initial points needed for finding the contacts and is + more robust for very small or large geom sizes. 2. Added :ref:`frame` to MJCF, a :ref:`meta-element` which defines a pure coordinate transformation on its direct children, without requiring a :ref:`body`. 3. Added the :at:`kv` attribute to the :ref:`position` and :ref:`intvelocity` @@ -64,17 +74,14 @@ General Plugins ^^^^^^^ - 4. Allow actuator plugins to use activation variables in ``mjData.act`` as their internal state, rather than ``mjData.plugin_state``. Actuator plugins can now specify :ref:`callbacks` that compute activation variables, and they can be used with built-in :ref:`dyntype` actuator dynamics. - 5. Added the `pid `__ actuator plugin, a configurable PID controller that implements the Integral term, which is not available with native MuJoCo actuators. MJX ^^^ - 6. Added ``site_xpos`` and ``site_xmat`` to MJX. 7. Added ``put_data``, ``put_model``, ``get_data`` to replace ``device_put`` and ``device_get_into``, which will be deprecated. These new functions correctly translate fields that are the result of intermediate calculations such as @@ -86,8 +93,8 @@ Bug fixes Before this fix such actuators could lead to non-conservation of momentum. 9. Fix bug that prevented using flex with :ref:`simulate`. 10. Fix bug that prevented the use of elasticity plugins in combination with pinned flex vertices. -11. Release Python wheels targeting macOS 10.16 to support x86_64 systems where SYSTEM_VERSION_COMPAT is set. The minimum - supported version is still 11.0, but we release these wheels to fix compatibility for those users. See +11. Release Python wheels targeting macOS 10.16 to support x86_64 systems where ``SYSTEM_VERSION_COMPAT`` is set. + The minimum supported version is still 11.0, but we release these wheels to fix compatibility for those users. See :github:issue:`1213`. Version 3.0.1 (November 15, 2023) diff --git a/doc/python.rst b/doc/python.rst index e18149cb..6f1ce280 100644 --- a/doc/python.rst +++ b/doc/python.rst @@ -465,28 +465,32 @@ Open-loop rollouts ================== We include a code sample showing how to add additional C/C++ functionality, exposed as a Python module via pybind11. The -sample, implemented in ``rollout.cc`` and wrapped in ``rollout.py``, implements a common use case where tight loops -implemented outside of Python are beneficial: rolling out a trajectory (i.e., calling ``mj_step()`` in a loop), given an -intial state and sequence of controls, and returning subsequent states and sensor values. The canonical usage form is +sample, implemented in `rollout.cc `__ +and wrapped in `rollout.py `__, +implements a common use case where tight loops implemented outside of Python are beneficial: rolling out a trajectory +(i.e., calling ``mj_step()`` in a loop), given an intial state and sequence of controls, and returning subsequent states +and sensor values. The basic usage form is .. code-block:: python - state, sensordata = rollout.rollout(model, data, initial_state, ctrl) + state, sensordata = rollout.rollout(model, data, initial_state, control) -``initial_state`` is a ``nstate x nqva`` array, with ``nstate`` initial states of length ``nqva``, where ``nqva = -model.nq + model.nv + model.na`` is the size of the full MuJoCo mechanical state: positions (``data.qpos``), velocities -(``data.qvel``) and actuator activations (``data.act``). ``ctrl`` is a ``nstate x nstep x nu`` array of control -sequences. +``initial_state`` is a ``nroll x nstate`` array, with ``nroll`` initial states of size ``nstate``, where +``nstate = mj_stateSize(model, mjtState.mjSTATE_FULLPHYSICS)`` is the size of the +:ref:`full physics state`. ``control`` is a ``nroll x nstep x ncontrol`` array of controls. Controls are +by default the ``mjModel.nu`` standard actuators, but any combination of :ref:`user input` arrays can be +specified by passing an optional ``control_spec`` bitflag. + +If a rollout diverges, the current state and sensor values are used to fill the remainder of the trajectory. +Therefore, non-increasing time values can be used to detect diverged rollouts. The ``rollout`` function is designed to be completely stateless, so all inputs of the stepping pipeline are set and any -values already present in the given ``MjData`` instance will have no effect on the output. In order to facilitate this, -all inputs including ``time`` and ``qacc_warmstart`` are set to default values, as are auxillary controls -(``qfrc_applied``, ``xfrc_applied`` and ``mocap_{pos,quat}``). These can also be optionally set by the user. +values already present in the given ``MjData`` instance will have no effect on the output. Since the Global Interpreter Lock can be released, this function can be efficiently threaded using Python threads. See the ``test_threading`` function in -`rollout_test.py `_ for an example of -threaded operation. +`rollout_test.py `__ for an example +of threaded operation (and more generally for usage examples). .. _PyMjpy_migration: diff --git a/python/mujoco/rollout.cc b/python/mujoco/rollout.cc index eacdc8fe..6cfbb470 100644 --- a/python/mujoco/rollout.cc +++ b/python/mujoco/rollout.cc @@ -31,111 +31,114 @@ namespace { namespace py = ::pybind11; +// NOLINTBEGIN(whitespace/line_length) + const auto rollout_doc = R"( -Roll out open-loop trajectories from initial states, get subsequent states and sensor values. +Roll out open-loop trajectories from initial states, get resulting states and sensor values. input arguments (required): - model an instance of MjModel - data an associated instance of MjData - nstate an integer, number of initial states from which to roll out trajectories - nstep an integer, number of steps to be taken for each trajectory + model instance of MjModel + data associated instance of MjData + nroll integer, number of initial states from which to roll out trajectories + nstep integer, number of steps to be taken for each trajectory + control_spec specification of controls, ncontrol = mj_stateSize(m, control_spec) + state0 (nroll x nstate) nroll initial state vectors, + nstate = mj_stateSize(m, mjSTATE_FULLPHYSICS) input arguments (optional): - initial_state (nstate x nqva) nstate initial state vectors, nqva=nq+nv+na - initial_time (nstate x 1) nstate initial times - initial_warmstart (nstate x nv) nstate qacc_warmstart vectors - ctrl (nstate x nstep x nu) nstate length-nstep controls - qfrc_applied (nstate x nstep x nv) nstate length-nstep generalized forces - xfrc_applied (nstate x nstep x nbody*6) nstate length-nstep Cartesian wrenches - mocap (nstate x nstep x nmocap*7) nstate length-nstep mocap body poses + warmstart0 (nroll x nv) nroll qacc_warmstart vectors + control (nroll x nstep x ncontrol) nroll trajectories of nstep controls output arguments (optional): - state (nstate x nstep x nqva) nstate length-nstep states - sensordata (nstate x nstep x nsendordata) nstate length-nstep sensordatas + state (nroll x nstep x nstate) nroll nstep states + sensordata (nroll x nstep x nsendordata) nroll trajectories of nstep sensordata vectors )"; // C-style rollout function, assumes all arguments are valid // all input fields of d are initialised, contents at call time do not matter // after returning, d will contain the last step of the last rollout -void _unsafe_rollout(const mjModel* m, mjData* d, int nstate, int nstep, - const mjtNum* state0, const mjtNum* ctrl, - const mjtNum* qfrc, const mjtNum* xfrc, - const mjtNum* mocap, const mjtNum* time0, - const mjtNum* warmstart0, +void _unsafe_rollout(const mjModel* m, mjData* d, int nroll, int nstep, unsigned int control_spec, + const mjtNum* state0, const mjtNum* warmstart0, const mjtNum* control, mjtNum* state, mjtNum* sensordata) { - // model sizes - int nq = m->nq; - int nv = m->nv; - int na = m->na; - int nqva = nq + nv + na; - int nu = m->nu; - int nbody = m->nbody; - int nmocap = m->nmocap; + // sizes + int nstate = mj_stateSize(m, mjSTATE_FULLPHYSICS); + int ncontrol = mj_stateSize(m, control_spec); + int nv = m->nv, nbody = m->nbody, neq = m->neq; int nsensordata = m->nsensordata; - // loop over initial states - for (int s=0; s < nstate; s++) { - // set initial state - if (state0) { - mju_copy(d->qpos, state0 + s*nqva, nq); - mju_copy(d->qvel, state0 + s*nqva + nq, nv); - mju_copy(d->act, state0 + s*nqva + nq + nv, na); - } else { - mju_copy(d->qpos, m->qpos0, nq); - mju_zero(d->qvel, nv); - mju_zero(d->act, na); + // clear user inputs if unspecified + if (!(control_spec & mjSTATE_CTRL)) { + mju_zero(d->ctrl, m->nu); + } + if (!(control_spec & mjSTATE_QFRC_APPLIED)) { + mju_zero(d->qfrc_applied, nv); + } + if (!(control_spec & mjSTATE_XFRC_APPLIED)) { + mju_zero(d->xfrc_applied, 6*nbody); + } + if (!(control_spec & mjSTATE_MOCAP_POS)) { + for (int i = 0; i < nbody; i++) { + int id = m->body_mocapid[i]; + if (id >= 0) mju_copy3(d->mocap_pos+3*id, m->body_pos+3*i); } + } + if (!(control_spec & mjSTATE_MOCAP_QUAT)) { + for (int i = 0; i < nbody; i++) { + int id = m->body_mocapid[i]; + if (id >= 0) mju_copy4(d->mocap_quat+4*id, m->body_quat+4*i); + } + } + if (!(control_spec & mjSTATE_EQ_ACTIVE)) { + for (int i = 0; i < neq; i++) { + d->eq_active[i] = m->eq_active0[i]; + } + } - // set initial time - d->time = time0 ? time0[s] : 0; + // loop over rollouts + for (int r = 0; r < nroll; r++) { + // set initial state + mj_setState(m, d, state0 + r*nstate, mjSTATE_FULLPHYSICS); // set warmstart accelerations if (warmstart0) { - mju_copy(d->qacc_warmstart, warmstart0 + s*nv, nv); + mju_copy(d->qacc_warmstart, warmstart0 + r*nv, nv); } else { mju_zero(d->qacc_warmstart, nv); } - // clear control inputs if unspecified - if (s == 0) { - if (!ctrl) { - mju_zero(d->ctrl, nu); - } - if (!qfrc) { - mju_zero(d->qfrc_applied, nv); - } - if (!xfrc) { - mju_zero(d->xfrc_applied, 6*nbody); - } - if (!mocap) { - for (int j=0; j < nbody; j++) { - int id = m->body_mocapid[j]; - if (id >= 0) { - mju_copy3(d->mocap_pos+3*id, m->body_pos+3*j); - mju_copy4(d->mocap_quat+4*id, m->body_quat+4*j); - } - } - } + // clear warning counters + for (int i = 0; i < mjNWARNING; i++) { + d->warning[i].number = 0; } - // roll out trajectories + // roll out trajectory for (int t = 0; t < nstep; t++) { + // check for warnings + bool nwarning = false; + for (int i = 0; i < mjNWARNING; i++) { + if (d->warning[i].number) { + nwarning = true; + break; + } + } + + // if any warnings, fill remaining outputs with current outputs, break + if (nwarning) { + for (; t < nstep; t++) { + int step = r*nstep + t; + if (state) { + mj_getState(m, d, state + step*nstate, mjSTATE_FULLPHYSICS); + } + if (sensordata) { + mju_copy(sensordata + step*nsensordata, d->sensordata, nsensordata); + } + } + break; + } + + int step = r*nstep + t; + // controls - if (ctrl) { - mju_copy(d->ctrl, ctrl + s*nstep*nu + t*nu, nu); - } - // generalized forces - if (qfrc) { - mju_copy(d->qfrc_applied, qfrc + s*nstep*nv + t*nv, nv); - } - // Cartesian wrenches - if (xfrc) { - mju_copy(d->xfrc_applied, xfrc + s*nstep*6*nbody + t*6*nbody, 6*nbody); - } - // mocap bodies - if (mocap) { - mju_copy(d->mocap_pos, - mocap + s*nstep*7*nmocap + t*7*nmocap, 3*nmocap); - mju_copy(d->mocap_quat, - mocap + s*nstep*7*nmocap + t*7*nmocap + 3*nmocap, 4*nmocap); + if (control) { + mj_setState(m, d, control + step*ncontrol, control_spec); } // step @@ -143,23 +146,22 @@ void _unsafe_rollout(const mjModel* m, mjData* d, int nstate, int nstep, // copy out new state if (state) { - mju_copy(state + s*nstep*nqva + t*nqva, d->qpos, nq); - mju_copy(state + s*nstep*nqva + t*nqva + nq, d->qvel, nv); - mju_copy(state + s*nstep*nqva + t*nqva + nq + nv, d->act, na); + mj_getState(m, d, state + step*nstate, mjSTATE_FULLPHYSICS); } + // copy out sensor values if (sensordata) { - mju_copy(sensordata + s*nstep*nsensordata + t*nsensordata, - d->sensordata, nsensordata); + mju_copy(sensordata + step*nsensordata, d->sensordata, nsensordata); } } } } +// NOLINTEND(whitespace/line_length) // check size of optional argument to rollout(), return raw pointer mjtNum* get_array_ptr(std::optional> arg, - const char* name, int nstate, int nstep, int dim) { + const char* name, int nroll, int nstep, int dim) { // if empty return nullptr if (!arg.has_value()) { return nullptr; @@ -169,11 +171,10 @@ mjtNum* get_array_ptr(std::optional> arg, py::buffer_info info = arg->request(); // check size - int expected_size = nstate * nstep * dim; + int expected_size = nroll * nstep * dim; if (info.size != expected_size) { std::ostringstream msg; - msg << name << ".size should be " << expected_size << - ", got " << info.size; + msg << name << ".size should be " << expected_size << ", got " << info.size; throw py::value_error(msg.str()); } return static_cast(info.ptr); @@ -188,14 +189,11 @@ PYBIND11_MODULE(_rollout, pymodule) { // get subsequent states and corresponding sensor values pymodule.def( "rollout", - [](const MjModelWrapper& m, MjDataWrapper& d, int nstate, int nstep, - std::optional init_state, - std::optional init_time, - std::optional init_warmstart, - std::optional ctrl, - std::optional qfrc, - std::optional xfrc, - std::optional mocap, + [](const MjModelWrapper& m, MjDataWrapper& d, + int nroll, int nstep, unsigned int control_spec, + const PyCArray state0, + std::optional warmstart0, + std::optional control, std::optional state, std::optional sensordata ) { @@ -203,28 +201,22 @@ PYBIND11_MODULE(_rollout, pymodule) { raw::MjData* data = d.get(); // check that some steps need to be taken, return if not - if (nstate < 1 || nstep < 1) { + if (nroll < 1 || nstep < 1) { return; } + // get sizes + int nstate = mj_stateSize(model, mjSTATE_FULLPHYSICS); + int ncontrol = mj_stateSize(model, control_spec); + // get raw pointers - int nqva = model->nq + model->nv + model->na; - mjtNum* init_state_ptr = - get_array_ptr(init_state, "initial_state", nstate, 1, nqva); - mjtNum* ctrl_ptr = - get_array_ptr(ctrl, "ctrl", nstate, nstep, model->nu); - mjtNum* qfrc_ptr = - get_array_ptr(qfrc, "qfrc_applied", nstate, nstep, model->nv); - mjtNum* xfrc_ptr = - get_array_ptr(xfrc, "xfrc_applied", nstate, nstep, 6*model->nbody); - mjtNum* mocap_ptr = - get_array_ptr(mocap, "mocap", nstate, nstep, 7*model->nmocap); - mjtNum* init_time_ptr = - get_array_ptr(init_time, "init_time", nstate, 1, 1); - mjtNum* init_warmstart_ptr = get_array_ptr( - init_warmstart, "init_warmstart", nstate, 1, model->nv); - mjtNum* state_ptr = get_array_ptr(state, "state", nstate, nstep, nqva); - mjtNum* sensordata_ptr = get_array_ptr(sensordata, "sensordata", nstate, + mjtNum* state0_ptr = get_array_ptr(state0, "state0", nroll, 1, nstate); + mjtNum* warmstart0_ptr = get_array_ptr(warmstart0, "warmstart0", nroll, + 1, model->nv); + mjtNum* control_ptr = get_array_ptr(control, "control", nroll, + nstep, ncontrol); + mjtNum* state_ptr = get_array_ptr(state, "state", nroll, nstep, nstate); + mjtNum* sensordata_ptr = get_array_ptr(sensordata, "sensordata", nroll, nstep, model->nsensordata); // perform rollouts @@ -234,24 +226,20 @@ PYBIND11_MODULE(_rollout, pymodule) { // call unsafe rollout function InterceptMjErrors(_unsafe_rollout)( - model, data, nstate, nstep, init_state_ptr, ctrl_ptr, qfrc_ptr, - xfrc_ptr, mocap_ptr, init_time_ptr, init_warmstart_ptr, state_ptr, - sensordata_ptr); + model, data, nroll, nstep, control_spec, state0_ptr, + warmstart0_ptr, control_ptr, state_ptr, sensordata_ptr); } }, py::arg("model"), py::arg("data"), - py::arg("nstate"), + py::arg("nroll"), py::arg("nstep"), - py::arg("initial_state") = py::none(), - py::arg("initial_time") = py::none(), - py::arg("initial_warmstart") = py::none(), - py::arg("ctrl") = py::none(), - py::arg("qfrc_applied") = py::none(), - py::arg("xfrc_applied") = py::none(), - py::arg("mocap") = py::none(), - py::arg("state") = py::none(), - py::arg("sensordata") = py::none(), + py::arg("control_spec"), + py::arg("state0"), + py::arg("warmstart0") = py::none(), + py::arg("control") = py::none(), + py::arg("state") = py::none(), + py::arg("sensordata") = py::none(), py::doc(rollout_doc) ); } @@ -259,3 +247,4 @@ PYBIND11_MODULE(_rollout, pymodule) { } // namespace } // namespace mujoco::python + diff --git a/python/mujoco/rollout.py b/python/mujoco/rollout.py index 10ad328f..8904e30c 100644 --- a/python/mujoco/rollout.py +++ b/python/mujoco/rollout.py @@ -14,132 +14,144 @@ # ============================================================================== """Roll out open-loop trajectories from initial states, get subsequent states and sensor values.""" +from typing import Optional + +import mujoco from mujoco import _rollout import numpy as np +from numpy import typing as npt -def rollout(model, data, initial_state=None, ctrl=None, - *, # require following arguments to be named - skip_checks=False, - nstate=None, - nstep=None, - initial_time=None, - initial_warmstart=None, - qfrc_applied=None, - xfrc_applied=None, - mocap=None, - state=None, - sensordata=None): - """Roll out open-loop trajectories from initial states, get subsequent states and sensor values. +def rollout(model: mujoco.MjModel, + data: mujoco.MjData, + initial_state: npt.ArrayLike, + control: Optional[npt.ArrayLike] = None, + *, # require subsequent arguments to be named + control_spec: int = mujoco.mjtState.mjSTATE_CTRL.value, + skip_checks: bool = False, + nroll: Optional[int] = None, + nstep: Optional[int] = None, + initial_warmstart: Optional[npt.ArrayLike] = None, + state: Optional[npt.ArrayLike] = None, + sensordata: Optional[npt.ArrayLike] = None): + """Rolls out open-loop trajectories from initial states, get subsequent states and sensor values. - This function serves as a Python wrapper for the C++ functionality in - `rollout.cc`, please see documentation therein. This python funtion will - infer `nstate` and `nstep`, tile input arguments with singleton dimensions, - and allocate output arguments if none are given. + Python wrapper for rollout.cc, see documentation therein. + Infers nroll and nstep. + Tiles inputs with singleton dimensions. + Allocates outputs if none are given. + + Args: + model: An mjModel instance. + data: An associated mjData instance. + initial_state: Array of initial states from which to roll out trajectories. + ([nroll or 1] x nstate) + control: Open-loop controls array to apply during the rollouts. + ([nroll or 1] x [nstep or 1] x ncontrol) + control_spec: mjtState specification of control vectors. + skip_checks: Whether to skip internal shape and type checks. + nroll: Number of rollouts (inferred if unspecified). + nstep: Number of steps in rollouts (inferred if unspecified). + initial_warmstart: Initial qfrc_warmstart array (optional). + ([nroll or 1] x nv) + state: State output array (optional). + (nroll x nstep x nstate) + sensordata: Sensor data output array (optional). + (nroll x nstep x nsensordata) + + Returns: + state: + State output array, (nroll x nstep x nstate). + sensordata: + Sensor data output array, (nroll x nstep x nsensordata). + + Raises: + ValueError: bad shapes or sizes. """ - # don't infer nstate/nstep, don't support singleton expansion, don't allocate - # output arrays, just call rollout + # skip_checks shortcut: + # don't infer nroll/nstep + # don't support singleton expansion + # don't allocate output arrays + # just call rollout and return if skip_checks: - _rollout.rollout(model, data, nstate, nstep, initial_state, initial_time, - initial_warmstart, ctrl, qfrc_applied, xfrc_applied, mocap, - state, sensordata) + _rollout.rollout(model, data, nroll, nstep, control_spec, initial_state, + initial_warmstart, control, state, sensordata) return state, sensordata + # check control_spec + if control_spec & ~mujoco.mjtState.mjSTATE_USER.value: + raise ValueError('control_spec can only contain bits in mjSTATE_USER') + # check types - if nstate and not isinstance(nstate, int): - raise ValueError('nstate must be an integer') + if nroll and not isinstance(nroll, int): + raise ValueError('nroll must be an integer') if nstep and not isinstance(nstep, int): raise ValueError('nstep must be an integer') _check_must_be_numeric( initial_state=initial_state, - initial_time=initial_time, initial_warmstart=initial_warmstart, - ctrl=ctrl, - qfrc_applied=qfrc_applied, - xfrc_applied=xfrc_applied, - mocap=mocap, + control=control, state=state, sensordata=sensordata) # check number of dimensions _check_number_of_dimensions(2, initial_state=initial_state, - initial_time=initial_time, initial_warmstart=initial_warmstart) _check_number_of_dimensions(3, - ctrl=ctrl, - qfrc_applied=qfrc_applied, - xfrc_applied=xfrc_applied, - mocap=mocap, + control=control, state=state, sensordata=sensordata) # ensure 2D, make contiguous, row-major (C ordering) initial_state = _ensure_2d(initial_state) - initial_time = _ensure_2d(initial_time) initial_warmstart = _ensure_2d(initial_warmstart) # ensure 3D, make contiguous, row-major (C ordering) - ctrl = _ensure_3d(ctrl) - qfrc_applied = _ensure_3d(qfrc_applied) - xfrc_applied = _ensure_3d(xfrc_applied) - mocap = _ensure_3d(mocap) + control = _ensure_3d(control) state = _ensure_3d(state) sensordata = _ensure_3d(sensordata) # check trailing dimensions - _check_trailing_dimension(model.nq + model.nv + model.na, - initial_state=initial_state, state=state) - _check_trailing_dimension(1, initial_time=initial_time) - _check_trailing_dimension(model.nu, ctrl=ctrl) - _check_trailing_dimension(model.nv, qfrc_applied=qfrc_applied) - _check_trailing_dimension(model.nbody*6, xfrc_applied=xfrc_applied) - _check_trailing_dimension(model.nmocap*7, mocap=mocap) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS.value) + _check_trailing_dimension(nstate, initial_state=initial_state, state=state) + ncontrol = mujoco.mj_stateSize(model, control_spec) + _check_trailing_dimension(ncontrol, control=control) + _check_trailing_dimension(model.nv, initial_warmstart=initial_warmstart) _check_trailing_dimension(model.nsensordata, sensordata=sensordata) - # infer nstate, check for incompatibilities - nstate = _infer_dimension(0, nstate or 1, - initial_state=initial_state, - initial_time=initial_time, - initial_warmstart=initial_warmstart, - ctrl=ctrl, - qfrc_applied=qfrc_applied, - xfrc_applied=xfrc_applied, - mocap=mocap, - state=state, - sensordata=sensordata) + # infer nroll, check for incompatibilities + nroll = _infer_dimension(0, nroll or 1, + initial_state=initial_state, + initial_warmstart=initial_warmstart, + control=control, + state=state, + sensordata=sensordata) # infer nstep, check for incompatibilities nstep = _infer_dimension(1, nstep or 1, - ctrl=ctrl, - qfrc_applied=qfrc_applied, - xfrc_applied=xfrc_applied, - mocap=mocap, + control=control, state=state, sensordata=sensordata) # tile input arrays if required (singleton expansion) - initial_state = _tile_if_required(initial_state, nstate) - initial_time = _tile_if_required(initial_time, nstate) - initial_warmstart = _tile_if_required(initial_warmstart, nstate) - ctrl = _tile_if_required(ctrl, nstate, nstep) - qfrc_applied = _tile_if_required(qfrc_applied, nstate, nstep) - xfrc_applied = _tile_if_required(xfrc_applied, nstate, nstep) - mocap = _tile_if_required(mocap, nstate, nstep) + initial_state = _tile_if_required(initial_state, nroll) + initial_warmstart = _tile_if_required(initial_warmstart, nroll) + control = _tile_if_required(control, nroll, nstep) # allocate output if not provided if state is None: - state = np.empty((nstate, nstep, model.nq + model.nv + model.na)) + state = np.empty((nroll, nstep, nstate)) if sensordata is None: - sensordata = np.empty((nstate, nstep, model.nsensordata)) + sensordata = np.empty((nroll, nstep, model.nsensordata)) # call rollout - _rollout.rollout(model, data, nstate, nstep, initial_state, initial_time, - initial_warmstart, ctrl, qfrc_applied, xfrc_applied, mocap, - state, sensordata) + _rollout.rollout(model, data, nroll, nstep, control_spec, initial_state, + initial_warmstart, control, state, sensordata) + + # return outputs + return state, sensordata - # return squeezed outputs - return state.squeeze(), sensordata.squeeze() def _check_must_be_numeric(**kwargs): for key, value in kwargs.items(): @@ -148,6 +160,7 @@ def _check_must_be_numeric(**kwargs): if not isinstance(value, np.ndarray) and not isinstance(value, float): raise ValueError(f'{key} must be a numpy array or float') + def _check_number_of_dimensions(ndim, **kwargs): for key, value in kwargs.items(): if value is None: @@ -155,12 +168,16 @@ def _check_number_of_dimensions(ndim, **kwargs): if value.ndim > ndim: raise ValueError(f'{key} can have at most {ndim} dimensions') + def _check_trailing_dimension(dim, **kwargs): for key, value in kwargs.items(): if value is None: continue if value.shape[-1] != dim: - raise ValueError(f'trailing dimension of {key} must be {dim}, got {value.shape[-1]}') + raise ValueError( + f'trailing dimension of {key} must be {dim}, got {value.shape[-1]}' + ) + def _ensure_2d(arg): if arg is None: @@ -168,6 +185,7 @@ def _ensure_2d(arg): else: return np.ascontiguousarray(np.atleast_2d(arg), dtype=np.float64) + def _ensure_3d(arg): if arg is None: return None @@ -181,7 +199,22 @@ def _ensure_3d(arg): arg = arg[np.newaxis, ...] return np.ascontiguousarray(arg, dtype=np.float64) + def _infer_dimension(dim, value, **kwargs): + """Infers dimension `dim` given guess `value` from set of arrays. + + Args: + dim: Dimension to be inferred. + value: Initial guess of inferred value (1: unknown). + **kwargs: List of arrays which should all have the same size (or 1) + along dimension dim. + + Returns: + Inferred dimension. + + Raises: + ValueError: If mismatch between array shapes or initial guess. + """ for name, array in kwargs.items(): if array is None: continue @@ -190,10 +223,12 @@ def _infer_dimension(dim, value, **kwargs): value = array.shape[dim] elif array.shape[dim] != 1: raise ValueError( - f'dimension {dim} inferred as {value} but {name} has {array.shape[dim]}' + f'dimension {dim} inferred as {value} ' + f'but {name} has {array.shape[dim]}' ) return value + def _tile_if_required(array, dim0, dim1=None): if array is None: return diff --git a/python/mujoco/rollout_test.py b/python/mujoco/rollout_test.py index 50ca4568..e1915592 100644 --- a/python/mujoco/rollout_test.py +++ b/python/mujoco/rollout_test.py @@ -14,15 +14,16 @@ # ============================================================================== """tests for rollout function.""" +import concurrent.futures +import threading + from absl.testing import absltest from absl.testing import parameterized import mujoco -import numpy as np -import concurrent.futures -import threading from mujoco import rollout +import numpy as np -#--------------------------- models used for testing --------------------------- +# -------------------------- models used for testing --------------------------- TEST_XML = r""" @@ -96,7 +97,7 @@ TEST_XML_MOCAP = r""" - + """ @@ -106,12 +107,33 @@ TEST_XML_EMPTY = r""" """ +TEST_XML_DIVERGE = r""" + + + + + + + + + + + + + + + +""" + ALL_MODELS = {'TEST_XML': TEST_XML, 'TEST_XML_NO_SENSORS': TEST_XML_NO_SENSORS, 'TEST_XML_NO_ACTUATORS': TEST_XML_NO_ACTUATORS, 'TEST_XML_EMPTY': TEST_XML_EMPTY} -#------------------------------- tests ----------------------------------------- +# ------------------------------ tests ----------------------------------------- + class MuJoCoRolloutTest(parameterized.TestCase): @@ -119,179 +141,209 @@ class MuJoCoRolloutTest(parameterized.TestCase): super().setUp() np.random.seed(42) - #----------------------------- test basic operation + # ----------------------------- test basic operation @parameterized.parameters(ALL_MODELS.keys()) def test_single_step(self, model_name): model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name]) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - initial_state = np.random.randn(model.nq + model.nv + model.na) - ctrl = np.random.randn(model.nu) - state, sensordata = rollout.rollout(model, data, initial_state, ctrl) + initial_state = np.random.randn(nstate) + control = np.random.randn(model.nu) + state, sensordata = rollout.rollout(model, data, initial_state, control) mujoco.mj_resetData(model, data) - py_state, py_sensordata = step(model, data, initial_state, ctrl=ctrl) + py_state, py_sensordata = py_rollout(model, data, initial_state, control) np.testing.assert_array_equal(state, py_state) np.testing.assert_array_equal(sensordata, py_sensordata) - - @parameterized.parameters(ALL_MODELS.keys()) - def test_single_rollout(self, model_name): + def test_one_rollout(self, model_name): nstep = 3 # number of timesteps model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name]) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - initial_state = np.random.randn(model.nq + model.nv + model.na) - ctrl = np.random.randn(nstep, model.nu) - state, sensordata = rollout.rollout(model, data, initial_state, ctrl) + initial_state = np.random.randn(nstate) + control = np.random.randn(nstep, model.nu) + state, sensordata = rollout.rollout(model, data, initial_state, control) - py_state, py_sensordata = single_rollout(model, data, initial_state, - ctrl=ctrl) - np.testing.assert_array_equal(state, np.asarray(py_state)) - np.testing.assert_array_equal(sensordata, np.asarray(py_sensordata)) + py_state, py_sensordata = py_rollout(model, data, initial_state, control) + np.testing.assert_array_equal(state, py_state) + np.testing.assert_array_equal(sensordata, py_sensordata) @parameterized.parameters(ALL_MODELS.keys()) def test_multi_step(self, model_name): model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name]) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - nstate = 5 # number of initial states + nroll = 5 # number of rollouts + nstep = 1 # number of steps - initial_state = np.random.randn(nstate, model.nq + model.nv + model.na) - ctrl = np.random.randn(nstate, 1, model.nu) - state, sensordata = rollout.rollout(model, data, initial_state, ctrl) + initial_state = np.random.randn(nroll, nstate) + control = np.random.randn(nroll, nstep, model.nu) + state, sensordata = rollout.rollout(model, data, initial_state, control) mujoco.mj_resetData(model, data) - py_state, py_sensordata = multi_rollout(model, data, initial_state, - ctrl=ctrl) + py_state, py_sensordata = py_rollout(model, data, initial_state, control) np.testing.assert_array_equal(state, py_state) np.testing.assert_array_equal(sensordata, py_sensordata) @parameterized.parameters(ALL_MODELS.keys()) - def test_single_rollout_fixed_ctrl(self, model_name): - nstep = 3 + def test_one_rollout_fixed_ctrl(self, model_name): model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name]) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - initial_state = np.random.randn(model.nq + model.nv + model.na) - ctrl = np.random.randn(model.nu) - state = np.empty((nstep, model.nq + model.nv + model.na)) - sensordata = np.empty((nstep, model.nsensordata)) - rollout.rollout(model, data, initial_state, ctrl, + nroll = 1 # number of rollouts + nstep = 3 # number of steps + + initial_state = np.random.randn(nstate) + control = np.random.randn(model.nu) + state = np.empty((nroll, nstep, nstate)) + sensordata = np.empty((nroll, nstep, model.nsensordata)) + rollout.rollout(model, data, initial_state, control, state=state, sensordata=sensordata) - ctrl = np.tile(ctrl, (nstep, 1)) # repeat?? - py_state, py_sensordata = single_rollout(model, data, initial_state, - ctrl=ctrl) + control = np.tile(control, (nstep, 1)) + py_state, py_sensordata = py_rollout(model, data, initial_state, control) np.testing.assert_array_equal(state, py_state) np.testing.assert_array_equal(sensordata, py_sensordata) @parameterized.parameters(ALL_MODELS.keys()) def test_multi_rollout(self, model_name): model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name]) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - nstate = 2 # number of initial states + nroll = 2 # number of initial states nstep = 3 # number of timesteps - initial_state = np.random.randn(nstate, model.nq + model.nv + model.na) - ctrl = np.random.randn(nstate, nstep, model.nu) - state, sensordata = rollout.rollout(model, data, initial_state, ctrl) + initial_state = np.random.randn(nroll, nstate) + control = np.random.randn(nroll, nstep, model.nu) + state, sensordata = rollout.rollout(model, data, initial_state, control) - py_state, py_sensordata = multi_rollout(model, data, initial_state, - ctrl=ctrl) - np.testing.assert_array_equal(py_state, py_state) - np.testing.assert_array_equal(py_sensordata, py_sensordata) + py_state, py_sensordata = py_rollout(model, data, initial_state, control) + np.testing.assert_array_equal(state, py_state) + np.testing.assert_array_equal(sensordata, py_sensordata) @parameterized.parameters(ALL_MODELS.keys()) def test_multi_rollout_fixed_ctrl_infer_from_output(self, model_name): model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name]) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - nstate = 2 # number of initial states + nroll = 2 # number of rollouts nstep = 3 # number of timesteps - initial_state = np.random.randn(nstate, model.nq + model.nv + model.na) - ctrl = np.random.randn(nstate, 1, model.nu) # 1 control in the time dimension - state = np.empty((nstate, nstep, model.nq + model.nv + model.na)) - state, sensordata = rollout.rollout(model, data, initial_state, ctrl, + initial_state = np.random.randn(nroll, nstate) + control = np.random.randn(nroll, 1, model.nu) + state = np.empty((nroll, nstep, nstate)) + state, sensordata = rollout.rollout(model, data, initial_state, control, state=state) - ctrl = np.repeat(ctrl, nstep, axis=1) - py_state, py_sensordata = multi_rollout(model, data, initial_state, - ctrl=ctrl) + control = np.repeat(control, nstep, axis=1) + py_state, py_sensordata = py_rollout(model, data, initial_state, control) np.testing.assert_array_equal(state, py_state) np.testing.assert_array_equal(sensordata, py_sensordata) - @parameterized.product(arg_nstep=[[3, 1, 1], [3, 3, 1], [3, 1, 3]], - model_name=list(ALL_MODELS.keys())) - def test_multi_rollout_multiple_inputs(self, arg_nstep, model_name): + @parameterized.parameters(ALL_MODELS.keys()) + def test_py_rollout_generalized_control(self, model_name): model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name]) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - nstate = 4 # number of initial states + nroll = 4 # number of rollouts + nstep = 3 # number of timesteps - initial_state = np.random.randn(nstate, model.nq + model.nv + model.na) + initial_state = np.random.randn(nroll, nstate) - # arg_nstep is the horizon for {ctrl, qfrc_applied, xfrc_applied}, respectively - ctrl = np.random.randn(nstate, arg_nstep[0], model.nu) - qfrc_applied = np.random.randn(nstate, arg_nstep[1], model.nv) - xfrc_applied = np.random.randn(nstate, arg_nstep[2], model.nbody*6) + control_spec = (mujoco.mjtState.mjSTATE_CTRL | + mujoco.mjtState.mjSTATE_QFRC_APPLIED | + mujoco.mjtState.mjSTATE_XFRC_APPLIED) + ncontrol = mujoco.mj_stateSize(model, control_spec) + control = np.random.randn(nroll, nstep, ncontrol) - state, sensordata = rollout.rollout(model, data, initial_state, ctrl, - qfrc_applied=qfrc_applied, - xfrc_applied=xfrc_applied) + state, sensordata = rollout.rollout(model, data, initial_state, control, + control_spec=control_spec) - # tile singleton arguments - nstep = max(arg_nstep) - if arg_nstep[0] == 1: - ctrl = np.repeat(ctrl, nstep, axis=1) - if arg_nstep[1] == 1: - qfrc_applied = np.repeat(qfrc_applied, nstep, axis=1) - if arg_nstep[2] == 1: - xfrc_applied = np.repeat(xfrc_applied, nstep, axis=1) - - py_state, py_sensordata = multi_rollout(model, data, initial_state, - ctrl=ctrl, - qfrc_applied=qfrc_applied, - xfrc_applied=xfrc_applied) + py_state, py_sensordata = py_rollout(model, data, initial_state, control, + control_spec=control_spec) np.testing.assert_array_equal(state, py_state) np.testing.assert_array_equal(sensordata, py_sensordata) - #----------------------------- test threaded operation + def test_detect_divergence(self): + model = mujoco.MjModel.from_xml_string(TEST_XML_DIVERGE) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) + data = mujoco.MjData(model) + + nroll = 4 # number of rollouts + initial_state = np.empty((nroll, nstate)) + + # get diverging (0, 2) and non-diverging (1, 3) states + mujoco.mj_getState(model, data, initial_state[0], + mujoco.mjtState.mjSTATE_FULLPHYSICS) + mujoco.mj_getState(model, data, initial_state[2], + mujoco.mjtState.mjSTATE_FULLPHYSICS) + mujoco.mj_resetDataKeyframe(model, data, 0) # keyframe 0 does not diverge + mujoco.mj_getState(model, data, initial_state[1], + mujoco.mjtState.mjSTATE_FULLPHYSICS) + mujoco.mj_getState(model, data, initial_state[3], + mujoco.mjtState.mjSTATE_FULLPHYSICS) + + nstep = 10000 # divergence after ~15s, timestep = 2e-3 + + state = np.random.randn(nroll, nstep, nstate) + + rollout.rollout(model, data, initial_state, state=state) + + # initial_state[0,2] diverged, final timesteps are identical + assert state[0][-1][0] == state[0][-2][0] + assert state[2][-1][0] == state[2][-2][0] + + # initial_state[1,3] did not diverge, final timesteps are different + assert state[1][-1][0] != state[1][-2][0] + assert state[3][-1][0] != state[3][-2][0] + + # ----------------------------- test threaded operation def test_threading(self): model = mujoco.MjModel.from_xml_string(TEST_XML) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) num_workers = 32 - nstate = 10000 + nroll = 10000 nstep = 5 - initial_state = np.random.randn(nstate, model.nq+model.nv+model.na) - state = np.zeros((nstate, nstep, model.nq+model.nv+model.na)) - sensordata = np.zeros((nstate, nstep, model.nsensordata)) - ctrl = np.random.randn(nstate, nstep, model.nu) + initial_state = np.random.randn(nroll, nstate) + state = np.empty((nroll, nstep, nstate)) + sensordata = np.empty((nroll, nstep, model.nsensordata)) + control = np.random.randn(nroll, nstep, model.nu) thread_local = threading.local() def thread_initializer(): thread_local.data = mujoco.MjData(model) - def call_rollout(initial_state, ctrl, state): - rollout.rollout(model, thread_local.data, skip_checks=True, - nstate=initial_state.shape[0], nstep=nstep, - initial_state=initial_state, ctrl=ctrl, state=state) + def call_rollout(initial_state, control, state, sensordata): + rollout.rollout(model, thread_local.data, initial_state, control, + skip_checks=True, nroll=initial_state.shape[0], + nstep=nstep, state=state, sensordata=sensordata) - n = initial_state.shape[0] // num_workers # integer division + n = nroll // num_workers # integer division chunks = [] # a list of tuples, one per worker for i in range(num_workers-1): - chunks.append( - (initial_state[i*n:(i+1)*n], ctrl[i*n:(i+1)*n], state[i*n:(i+1)*n])) + chunks.append((initial_state[i*n:(i+1)*n], + control[i*n:(i+1)*n], + state[i*n:(i+1)*n], + sensordata[i*n:(i+1)*n])) + # last chunk, absorbing the remainder: - chunks.append( - (initial_state[(num_workers-1)*n:], ctrl[(num_workers-1)*n:], - state[(num_workers-1)*n:])) + chunks.append((initial_state[(num_workers-1)*n:], + control[(num_workers-1)*n:], + state[(num_workers-1)*n:], + sensordata[(num_workers-1)*n:])) with concurrent.futures.ThreadPoolExecutor( max_workers=num_workers, initializer=thread_initializer) as executor: @@ -302,187 +354,237 @@ class MuJoCoRolloutTest(parameterized.TestCase): future.result() data = mujoco.MjData(model) - py_state, py_sensordata = multi_rollout(model, data, initial_state, - ctrl=ctrl) + py_state, py_sensordata = py_rollout(model, data, initial_state, control) np.testing.assert_array_equal(state, py_state) + np.testing.assert_array_equal(sensordata, py_sensordata) - #----------------------------- test advanced operation - - def test_time(self): - model = mujoco.MjModel.from_xml_string(TEST_XML) - data = mujoco.MjData(model) - - nstate = 1 - nstep = 3 - - initial_time = np.array([[2.]]) - initial_state = np.random.randn(nstate, model.nq + model.nv + model.na) - ctrl = np.random.randn(nstate, nstep, model.nu) - state, sensordata = rollout.rollout(model, data, initial_state, ctrl, - initial_time=initial_time) - - self.assertAlmostEqual(data.time, 2 + nstep*model.opt.timestep) + # ---------------------------- test advanced operation def test_warmstart(self): model = mujoco.MjModel.from_xml_string(TEST_XML) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - state0 = np.zeros(model.nq + model.nv + model.na) - ctrl = np.zeros(model.nu) - state1, _ = step(model, data, state0, ctrl=ctrl) + # take one step, save the state + state0 = np.zeros(nstate) + control = np.zeros(model.nu) + state1, _ = step(model, data, state0, control) + + # save qacc_warmstart initial_warmstart = data.qacc_warmstart.copy() - state2, _ = step(model, data, state1, ctrl=ctrl) + # take one more step (uses correct warmstart) + state2, _ = step(model, data, state1[0], control) - state, _ = rollout.rollout(model, data, state1, ctrl) - assert np.linalg.norm(state-state2) > 0 + # take step using rollout, don't take warmstart into account + state, _ = rollout.rollout(model, data, state1[0], control) - state, _ = rollout.rollout(model, data, state1, ctrl, + # assert that stepping without warmstarts is not exact + np.testing.assert_raises(AssertionError, + np.testing.assert_array_equal, state, state2) + + # take step using rollout, take warmstart into account + state, _ = rollout.rollout(model, data, state1, control, initial_warmstart=initial_warmstart) - np.testing.assert_array_equal(state, state2) + + # assert exact equality + np.testing.assert_array_equal(state, np.expand_dims(state2, axis=0)) def test_mocap(self): model = mujoco.MjModel.from_xml_string(TEST_XML_MOCAP) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - initial_state = np.zeros(model.nq + model.nv + model.na) + initial_state = np.zeros(nstate) + + control_spec = (mujoco.mjtState.mjSTATE_MOCAP_POS | + mujoco.mjtState.mjSTATE_MOCAP_QUAT) + pos1 = np.array((1., 2., 3.)) quat1 = np.array((1., 2., 3., 4.)) quat1 /= np.linalg.norm(quat1) pos2 = np.array((2., 3., 4.)) quat2 = np.array((2., 3., 4., 5.)) quat2 /= np.linalg.norm(quat2) - mocap = np.hstack((pos1, quat1, pos2, quat2)) + control = np.hstack((pos1, pos2, quat1, quat2)) - state, sensordata = rollout.rollout(model, data, initial_state, mocap=mocap) + _, sensordata = rollout.rollout(model, data, initial_state, control, + control_spec=control_spec) - np.testing.assert_array_almost_equal(sensordata[:3], pos1) - np.testing.assert_array_almost_equal(sensordata[3:], quat2) + np.testing.assert_array_almost_equal(sensordata[0][0][:3], pos1) + np.testing.assert_array_almost_equal(sensordata[0][0][3:], quat1) - #----------------------------- test correctness + # ---------------------------- test correctness def test_intercept_mj_errors(self): model = mujoco.MjModel.from_xml_string(TEST_XML) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - initial_state = np.zeros(model.nq + model.nv + model.na) - ctrl = np.zeros((3, model.nu)) + nroll = 1 + nstep = 3 + + initial_state = np.zeros((nroll, nstate)) + ctrl = np.zeros((nroll, nstep, model.nu)) model.opt.solver = 10 # invalid solver type with self.assertRaisesWithLiteralMatch( mujoco.FatalError, 'mj_fwdConstraint: unknown solver type 10'): - state, sensordata = rollout.rollout(model, data, initial_state, ctrl) + rollout.rollout(model, data, initial_state, ctrl) def test_invalid(self): model = mujoco.MjModel.from_xml_string(TEST_XML) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - initial_state = np.zeros(model.nq + model.nv + model.na) + nroll = 1 - ctrl = 'string' - with self.assertRaisesWithLiteralMatch( - ValueError, 'ctrl must be a numpy array or float'): - state, sensordata = rollout.rollout(model, data, initial_state, ctrl) + initial_state = np.zeros((nroll, nstate)) - qfrc_applied = np.zeros((2, 3, 4, 5)) + control = 'string' with self.assertRaisesWithLiteralMatch( - ValueError, 'qfrc_applied can have at most 3 dimensions'): - state, sensordata = rollout.rollout(model, data, initial_state, - qfrc_applied=qfrc_applied) + ValueError, 'control must be a numpy array or float'): + rollout.rollout(model, data, initial_state, control) + + control = np.zeros((2, 3, 4, 5)) + with self.assertRaisesWithLiteralMatch( + ValueError, 'control can have at most 3 dimensions'): + rollout.rollout(model, data, initial_state, control) def test_bad_sizes(self): model = mujoco.MjModel.from_xml_string(TEST_XML) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - initial_state = np.random.randn(model.nq + model.nv + model.na+1) - with self.assertRaisesWithLiteralMatch( - ValueError, 'trailing dimension of initial_state must be 5, got 6'): - state, sensordata = rollout.rollout(model, data, initial_state) + nroll = 1 + nstep = 3 - initial_state = np.random.randn(model.nq + model.nv + model.na) - ctrl = np.random.randn(model.nu+1) + initial_state = np.random.randn(nroll, nstate + 1) with self.assertRaisesWithLiteralMatch( - ValueError, 'trailing dimension of ctrl must be 2, got 3'): - state, sensordata = rollout.rollout(model, data, initial_state, ctrl) + ValueError, 'trailing dimension of initial_state must be 6, got 7'): + rollout.rollout(model, data, initial_state) - ctrl = np.random.randn(2, model.nu) - qfrc_applied = np.random.randn(3, model.nv) # incompatible horizon + initial_state = np.random.randn(nroll, nstate) + control = np.random.randn(1, nstep, model.nu + 1) with self.assertRaisesWithLiteralMatch( - ValueError, 'dimension 1 inferred as 2 but qfrc_applied has 3'): - state, sensordata = rollout.rollout(model, data, initial_state, ctrl, - qfrc_applied=qfrc_applied) + ValueError, 'trailing dimension of control must be 2, got 3'): + rollout.rollout(model, data, initial_state, control) + + control = np.random.randn(nroll, nstep, model.nu) + state = np.random.randn(nroll, nstep+1, nstate) # incompatible nstep + with self.assertRaisesWithLiteralMatch( + ValueError, 'dimension 1 inferred as 3 but state has 4'): + rollout.rollout(model, data, initial_state, control, state=state) + + initial_state = np.random.randn(nroll, nstate) + control = np.random.randn(nroll, nstep, model.nu) + bad_spec = mujoco.mjtState.mjSTATE_ACT + with self.assertRaisesWithLiteralMatch( + ValueError, 'control_spec can only contain bits in mjSTATE_USER'): + rollout.rollout(model, data, initial_state, control, + control_spec=bad_spec) def test_stateless(self): model = mujoco.MjModel.from_xml_string(TEST_XML) - model.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_WARMSTART.value + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - # call step with a clean mjData - initial_state = np.random.randn(model.nq + model.nv + model.na) - ctrl = np.random.randn(model.nu) - state, sensordata = rollout.rollout(model, data, initial_state, ctrl) + # step with a clean mjData + initial_state = np.random.randn(nstate) + control = np.random.randn(3, 3, model.nu) + state, sensordata = rollout.rollout(model, data, initial_state, control) - # fill mjData with some debug value, see that we still get the same outputs - mujoco.mj_resetDataDebug(model, data, 255) - debug_state, debug_sensordata = rollout.rollout(model, data, initial_state, - ctrl) + # fill user fields with random values + for attr in [ + 'ctrl', + 'qfrc_applied', + 'xfrc_applied', + 'mocap_pos', + 'mocap_quat', + ]: + setattr(data, attr, np.random.randn(*getattr(data, attr).shape)) - np.testing.assert_array_equal(state, debug_state) - np.testing.assert_array_equal(sensordata, debug_sensordata) + # roll out again + state2, sensordata2 = rollout.rollout(model, data, initial_state, control) + + # assert that we still get the same outputs + np.testing.assert_array_equal(state, state2) + np.testing.assert_array_equal(sensordata, sensordata2) -#--------------- Python implementation of rollout functionality ---------------- +# -------------- Python implementation of rollout functionality ---------------- -def get_state(data): - return np.hstack((data.qpos, data.qvel, data.act)) -def set_state(model, data, state): - data.qpos = state[:model.nq] - data.qvel = state[model.nq:model.nq+model.nv] - data.act = state[model.nq+model.nv:model.nq+model.nv+model.na] +def get_state(model, data): + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) + state = np.empty(nstate) + mujoco.mj_getState(model, data, state, mujoco.mjtState.mjSTATE_FULLPHYSICS) + return state.reshape((1, nstate)) -def step(model, data, state, **kwargs): + +def step(model, data, state, control, + control_spec=mujoco.mjtState.mjSTATE_CTRL): if state is not None: - set_state(model, data, state) - for key, value in kwargs.items(): - if value is not None: - setattr(data, key, np.reshape(value, getattr(data, key).shape)) + mujoco.mj_setState(model, data, state, mujoco.mjtState.mjSTATE_FULLPHYSICS) + mujoco.mj_setState(model, data, control, control_spec) mujoco.mj_step(model, data) - return (get_state(data), data.sensordata) + return (get_state(model, data), data.sensordata) -def single_rollout(model, data, initial_state, **kwargs): - arg_nstep = set([a.shape[0] for a in kwargs.values()]) - assert len(arg_nstep) == 1 # nstep dimensions must match - nstep = arg_nstep.pop() - state = np.empty((nstep, model.nq + model.nv + model.na)) +def one_rollout(model, data, initial_state, control, + control_spec=mujoco.mjtState.mjSTATE_CTRL): + nstep = control.shape[0] + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) + state = np.empty((nstep, nstate)) sensordata = np.empty((nstep, model.nsensordata)) mujoco.mj_resetData(model, data) for t in range(nstep): - kwargs_t = {} - for key, value in kwargs.items(): - kwargs_t[key] = value[0 if value.ndim == 1 else t] state[t], sensordata[t] = step(model, data, - initial_state if t==0 else None, - **kwargs_t) + initial_state if t == 0 else None, + control[t], control_spec) return state, sensordata -def multi_rollout(model, data, initial_state, **kwargs): - nstate = initial_state.shape[0] - arg_nstep = set([a.shape[1] for a in kwargs.values()]) - assert len(arg_nstep) == 1 # nstep dimensions must match - nstep = arg_nstep.pop() - state = np.empty((nstate, nstep, model.nq + model.nv + model.na)) - sensordata = np.empty((nstate, nstep, model.nsensordata)) - for s in range(nstate): - kwargs_s = {key : value[s] for key, value in kwargs.items()} - state_s, sensordata_s = single_rollout(model, data, initial_state[s], - **kwargs_s) - state[s] = state_s - sensordata[s] = sensordata_s - return state.squeeze(), sensordata.squeeze() +def ensure_2d(arg): + if arg is None: + return None + else: + return np.ascontiguousarray(np.atleast_2d(arg), dtype=np.float64) + + +def ensure_3d(arg): + if arg is None: + return None + else: + # np.atleast_3d adds both leading and trailing dims, we want only leading + if arg.ndim == 0: + arg = arg[np.newaxis, np.newaxis, np.newaxis, ...] + elif arg.ndim == 1: + arg = arg[np.newaxis, np.newaxis, ...] + elif arg.ndim == 2: + arg = arg[np.newaxis, ...] + return np.ascontiguousarray(arg, dtype=np.float64) + + +def py_rollout(model, data, initial_state, control, + control_spec=mujoco.mjtState.mjSTATE_CTRL): + initial_state = ensure_2d(initial_state) + control = ensure_3d(control) + nroll = initial_state.shape[0] + nstep = control.shape[1] + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) + + state = np.empty((nroll, nstep, nstate)) + sensordata = np.empty((nroll, nstep, model.nsensordata)) + for r in range(nroll): + state_r, sensordata_r = one_rollout( + model, data, initial_state[r], control[r], control_spec + ) + state[r] = state_r + sensordata[r] = sensordata_r + return state, sensordata + if __name__ == '__main__': absltest.main() From 0a7be1732ca8b1c6f4d5abb37dd1a179935cc2a2 Mon Sep 17 00:00:00 2001 From: Erik Frey Date: Tue, 23 Jan 2024 00:31:01 -0800 Subject: [PATCH 25/92] Adds support for explicit dense/sparse mass matrices to MJX. This increases performance, particularly for the Newton solver on TPU. PiperOrigin-RevId: 600696483 Change-Id: If69bb9a2e21ba8dad6ca23f093ce7b7ceae644ff --- doc/changelog.rst | 18 +++-- doc/mjx.rst | 8 ++ mjx/mujoco/mjx/__init__.py | 4 +- mjx/mujoco/mjx/_src/device_test.py | 2 + mjx/mujoco/mjx/_src/forward.py | 6 +- mjx/mujoco/mjx/_src/io.py | 44 ++++++++++- mjx/mujoco/mjx/_src/io_test.py | 34 ++++++--- mjx/mujoco/mjx/_src/smooth.py | 92 +++++------------------ mjx/mujoco/mjx/_src/smooth_test.py | 21 +----- mjx/mujoco/mjx/_src/solver.py | 19 ++--- mjx/mujoco/mjx/_src/solver_test.py | 40 ++++++++-- mjx/mujoco/mjx/_src/support.py | 109 +++++++++++++++++++++++++++- mjx/mujoco/mjx/_src/support_test.py | 40 ++++++++++ mjx/mujoco/mjx/_src/types.py | 30 ++++++-- 14 files changed, 324 insertions(+), 143 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index cb27d8b4..40b74e01 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -26,21 +26,23 @@ MJX 5. Added :at:`site` transmission. 6. Updated MJX colab tutorial with more stable quadruped environment. 7. Added ``mjx.ray`` which mirrors :ref:`mj_ray` for planes, spheres, capsules, boxes, and meshes. +8. Added ``mjx.is_sparse`` which mirrors :ref:`mj_isSparse` and ``mjx.full_m`` which mirrors :ref:`mj_fullM`. +9. Added support for specifying sparse or dense mass matrices via :ref:`option-jacobian`. Python bindings ^^^^^^^^^^^^^^^ -8. Improved the implmentation of the :ref:`rollout` module. Note the changes below are breaking, dependent - code will require modification. +10. Improved the implmentation of the :ref:`rollout` module. Note the changes below are breaking, dependent + code will require modification. - - Uses :ref:`mjSTATE_FULLPHYSICS` as state spec, enabling divergence detection by inspecting time. - - Allows user-defined control spec for any combination of :ref:`user input` fields as controls. - - Outputs are no longer squeezed and always have dim=3. + - Uses :ref:`mjSTATE_FULLPHYSICS` as state spec, enabling divergence detection by inspecting time. + - Allows user-defined control spec for any combination of :ref:`user input` fields as controls. + - Outputs are no longer squeezed and always have dim=3. Bug fixes ^^^^^^^^^ -9. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes - :github:issue:`1270`. -10. Fixed a bug in the :ref:`muscle model` that led to non-zero values outside the lower +11. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes + :github:issue:`1270`. +12. Fixed a bug in the :ref:`muscle model` that led to non-zero values outside the lower bound of the length range. Fixes :github:issue:`1342`. diff --git a/doc/mjx.rst b/doc/mjx.rst index 251e05fc..50c47389 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -349,3 +349,11 @@ For MJX to perform well, some configuration parameters should be adjusted from t :ref:`option-flag` element Disabling ``eulerdamp`` can help performance and is often not needed for stability. + +:ref:`option-jacobian` element + Explicitly setting "dense" or "sparse" may speed up simulation depending on your device. Modern TPUs have specialized + hardware for rapidly operating over sparse matrices, whereas GPUs tend to be faster with dense matrices as long as + they fit onto the device. As such, the behavior in MJX for the default "auto" setting is sparse if ``nv`` is 60 or + greater, or if MJX detects a TPU as the default backend, otherwise "dense". For TPU, using "sparse" with the + Newton solver can speed up simulation by 2x to 3x. For GPU, choosing "dense" may impart a more modest speedup of 10% + to 20%, as long as the dense matrices can fit on the device. diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py index f4953c60..737fe2c5 100644 --- a/mjx/mujoco/mjx/__init__.py +++ b/mjx/mujoco/mjx/__init__.py @@ -39,8 +39,10 @@ from mujoco.mjx._src.smooth import com_vel from mujoco.mjx._src.smooth import crb from mujoco.mjx._src.smooth import factor_m from mujoco.mjx._src.smooth import kinematics -from mujoco.mjx._src.smooth import mul_m from mujoco.mjx._src.smooth import rne from mujoco.mjx._src.smooth import transmission from mujoco.mjx._src.solver import solve +from mujoco.mjx._src.support import is_sparse +from mujoco.mjx._src.support import full_m +from mujoco.mjx._src.support import mul_m from mujoco.mjx._src.types import * diff --git a/mjx/mujoco/mjx/_src/device_test.py b/mjx/mujoco/mjx/_src/device_test.py index 71137c5d..6126f707 100644 --- a/mjx/mujoco/mjx/_src/device_test.py +++ b/mjx/mujoco/mjx/_src/device_test.py @@ -75,6 +75,7 @@ class DeviceTest(parameterized.TestCase): def testdevice_get(self, fname): """Test getting MjData from a device.""" m = test_util.load_test_file(fname) + m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE # force sparse for testing mx = device.device_put(m) dx = mjx.make_data(mx) d = mujoco.MjData(m) @@ -85,6 +86,7 @@ class DeviceTest(parameterized.TestCase): def testdevice_get_batched(self, fname): """Test getting MjData from a device.""" m = test_util.load_test_file(fname) + m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE # force sparse for testing mx = device.device_put(m) batch_size = 32 diff --git a/mjx/mujoco/mjx/_src/forward.py b/mjx/mujoco/mjx/_src/forward.py index 1ff268e7..682cc88e 100644 --- a/mjx/mujoco/mjx/_src/forward.py +++ b/mjx/mujoco/mjx/_src/forward.py @@ -66,7 +66,7 @@ def fwd_position(m: Model, d: Data) -> Data: d = smooth.kinematics(m, d) d = smooth.com_pos(m, d) d = smooth.crb(m, d) - d = smooth.factor_m(m, d, d.qM) + d = smooth.factor_m(m, d) d = collision_driver.collision(m, d) d = constraint.make_constraint(m, d) d = smooth.transmission(m, d) @@ -288,8 +288,8 @@ def euler(m: Model, d: Data) -> Data: qacc = d.qacc if not m.opt.disableflags & DisableBit.EULERDAMP: # TODO(robotics-simulation): can this be done with a smaller perf hit - mh = d.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping) - dh = smooth.factor_m(m, d, mh) + dh = d.replace(qM=d.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping)) + dh = smooth.factor_m(m, dh) qfrc = d.qfrc_smooth + d.qfrc_constraint qacc = smooth.solve_m(m, dh, qfrc) return _advance(m, d, d.act_dot, qacc) diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 5b406185..dd91de94 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -23,8 +23,10 @@ import mujoco from mujoco.mjx._src import collision_driver from mujoco.mjx._src import constraint from mujoco.mjx._src import mesh +from mujoco.mjx._src import support from mujoco.mjx._src import types import numpy as np +import scipy def _put_option(o: mujoco.MjOption, device=None) -> types.Option: @@ -35,6 +37,9 @@ def _put_option(o: mujoco.MjOption, device=None) -> types.Option: if o.cone not in set(types.ConeType): raise NotImplementedError(f'{mujoco.mjtCone(o.cone)}') + if o.jacobian not in set(types.JacobianType): + raise NotImplementedError(f'{mujoco.mjtJacobian(o.jacobian)}') + if o.solver not in set(types.SolverType): raise NotImplementedError(f'{mujoco.mjtSolver(o.solver)}') @@ -49,6 +54,7 @@ def _put_option(o: mujoco.MjOption, device=None) -> types.Option: } static_fields['integrator'] = types.IntegratorType(o.integrator) static_fields['cone'] = types.ConeType(o.cone) + static_fields['jacobian'] = types.JacobianType(o.jacobian) static_fields['solver'] = types.SolverType(o.solver) static_fields['disableflags'] = types.DisableBit(o.disableflags) @@ -137,8 +143,10 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data: ne, nf, nl, nc = constraint.count_constraints(m) nefc = ne + nf + nl + nc + zero_0 = jp.zeros(0, dtype=jp.float32) zero_nv = jp.zeros(m.nv, dtype=jp.float32) zero_nv_6 = jp.zeros((m.nv, 6), dtype=jp.float32) + zero_nv_nv = jp.zeros((m.nv, m.nv), dtype=jp.float32) zero_nbody_3 = jp.zeros((m.nbody, 3), dtype=jp.float32) zero_nbody_6 = jp.zeros((m.nbody, 6), dtype=jp.float32) zero_nbody_10 = jp.zeros((m.nbody, 10), dtype=jp.float32) @@ -180,10 +188,9 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data: actuator_length=zero_nu, actuator_moment=jp.zeros((m.nu, m.nv), dtype=jp.float32), crb=zero_nbody_10, - qM=zero_nm, - qLD=zero_nm, - qLDiagInv=zero_nv, - qLDiagSqrtInv=zero_nv, + qM=zero_nm if support.is_sparse(m) else zero_nv_nv, + qLD=zero_nm if support.is_sparse(m) else zero_nv_nv, + qLDiagInv=zero_nv if support.is_sparse(m) else zero_0, contact=types.Contact.zero(ncon), efc_J=jp.zeros((nefc, m.nv), dtype=jp.float32), efc_frictionloss=zero_nefc, @@ -237,6 +244,14 @@ def get_data( mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL, ]).repeat([ne, nf, nl, nc]) + dof_i, dof_j = [], [] + for i in range(m.nv): + j = i + while j > -1: + dof_i.append(i) + dof_j.append(j) + j = m.dof_parentid[j] + ds = [] for i in range(batch_size): dx_i = jax.tree_map(lambda x, i=i: x[i], dx) if batched else d @@ -267,6 +282,15 @@ def get_data( if field.name == 'efc_J': value = value[efc_active].reshape(-1) + if field.name == 'qM' and not support.is_sparse(m): + value = value[dof_i, dof_j] + + if field.name == 'qLD' and not support.is_sparse(m): + value = value[dof_i, dof_j] + + if field.name == 'qLDiagInv' and not support.is_sparse(m): + value = np.ones(m.nv) + if value.shape: getattr(d_i, field.name)[:] = value else: @@ -346,6 +370,18 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data: value[value_beg:value_beg+size] = fields[fname][d_beg:d_beg+size] fields[fname] = value + # convert qM and qLD if jacobian is dense + if not support.is_sparse(m): + fields['qM'] = np.zeros((m.nv, m.nv)) + mujoco.mj_fullM(m, fields['qM'], d.qM) + # TODO(erikfrey): derive L*L' from L'*D*L instead of recomputing + try: + fields['qLD'], _ = scipy.linalg.cho_factor(fields['qM']) + except scipy.linalg.LinAlgError: + # this happens when qM is empty or unstable simulation + fields['qLD'] = np.zeros((m.nv, m.nv)) + fields['qLDiagInv'] = np.zeros(0) + fields = jax.device_put(fields, device=device) fields['contact'] = _put_contact(d.contact, ncon, device=device) diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index 8b25f359..b402c407 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -25,7 +25,7 @@ import numpy as np _MULTIPLE_CONVEX_OBJECTS = """ - )"; mjModel* model = LoadModelFromString(xml); - ASSERT_THAT(model, testing::NotNull()); + ASSERT_THAT(model, NotNull()); mj_deleteModel(model); } @@ -437,7 +438,7 @@ TEST_F(MjCMeshTest, SmallInertiaLoads) { )"; mjModel* model = LoadModelFromString(xml); - ASSERT_THAT(model, testing::NotNull()); + ASSERT_THAT(model, NotNull()); mj_deleteModel(model); } @@ -526,7 +527,7 @@ TEST_F(MjCMeshTest, FlippedFaceAllowedWorld) { )"; std::array error; mjModel* model = LoadModelFromString(xml, error.data(), error.size()); - EXPECT_THAT(model, testing::NotNull()); + EXPECT_THAT(model, NotNull()); CheckTetrahedronWasRescaled(model); mj_deleteModel(model); } @@ -548,7 +549,7 @@ TEST_F(MjCMeshTest, FlippedFaceAllowedNoMass) { )"; std::array error; mjModel* model = LoadModelFromString(xml, error.data(), error.size()); - EXPECT_THAT(model, testing::NotNull()); + EXPECT_THAT(model, NotNull()); CheckTetrahedronWasRescaled(model); mj_deleteModel(model); } @@ -571,7 +572,7 @@ TEST_F(MjCMeshTest, FlippedFaceAllowedInertial) { )"; std::array error; mjModel* model = LoadModelFromString(xml, error.data(), error.size()); - EXPECT_THAT(model, testing::NotNull()); + EXPECT_THAT(model, NotNull()); CheckTetrahedronWasRescaled(model); mj_deleteModel(model); } @@ -593,7 +594,7 @@ TEST_F(MjCMeshTest, FlippedFaceAllowedNegligibleArea) { )"; std::array error; mjModel* model = LoadModelFromString(xml, error.data(), error.size()); - EXPECT_THAT(model, testing::NotNull()); + EXPECT_THAT(model, NotNull()); CheckTetrahedronWasRescaled(model); mj_deleteModel(model); } @@ -649,7 +650,7 @@ TEST_F(MjCMeshTest, AreaTooSmallAllowedWorld) { )"; std::array error; mjModel* model = LoadModelFromString(xml, error.data(), error.size()); - EXPECT_THAT(model, testing::NotNull()); + EXPECT_THAT(model, NotNull()); mj_deleteModel(model); } @@ -692,7 +693,7 @@ TEST_F(MjCMeshTest, VolumeSmallAllowedShell) { )"; std::array error; mjModel* model = LoadModelFromString(xml, error.data(), error.size()); - ASSERT_THAT(model, testing::NotNull()); + ASSERT_THAT(model, NotNull()); EXPECT_LE(mju_abs(model->geom_size[0]), 1); EXPECT_LE(mju_abs(model->geom_size[1]), 1); EXPECT_LE(mju_abs(model->geom_size[2]), 1); @@ -717,7 +718,7 @@ TEST_F(MjCMeshTest, VolumeNegativeDefaultsLegacy) { )"; std::array error; mjModel* model = LoadModelFromString(xml, error.data(), error.size()); - EXPECT_THAT(model, testing::NotNull()); + EXPECT_THAT(model, NotNull()); EXPECT_LE(mju_abs(model->geom_size[0]), 1); EXPECT_LE(mju_abs(model->geom_size[1]), 1); EXPECT_LE(mju_abs(model->geom_size[2]), 1); @@ -740,7 +741,7 @@ TEST_F(MjCMeshTest, VolumeTooSmallAllowedWorld) { )"; std::array error; mjModel* model = LoadModelFromString(xml, error.data(), error.size()); - EXPECT_THAT(model, testing::NotNull()); + EXPECT_THAT(model, NotNull()); mj_deleteModel(model); } @@ -818,7 +819,7 @@ TEST_F(MjCMeshTest, MeshPosQuat) { )"; mjModel* model = LoadModelFromString(xml); - ASSERT_THAT(model, testing::NotNull()); + ASSERT_THAT(model, NotNull()); // Loading the mesh results in an offset of the geom's pos and quat due to the // fact that the geom's center is not the volumetric center of the mesh. To // recover the geom's originally specified pose, the offset used is stored in @@ -862,5 +863,66 @@ TEST_F(MjCMeshTest, MeshPosQuat) { mj_deleteModel(model); } +// ----------------------------- texcoord ------------------------------------- + +TEST_F(MjCMeshTest, CreateFaceTexCoord) { + static constexpr char xml[] = R"( + + + + + + )"; + std::array error; + mjModel* model = LoadModelFromString(xml, error.data(), error.size()); + EXPECT_THAT(model, NotNull()) << error.data(); + mj_deleteModel(model); +} + +TEST_F(MjCMeshTest, UseFaceTexCoord) { + static constexpr char xml[] = R"( + + + + + + )"; + std::array error; + mjModel* model = LoadModelFromString(xml, error.data(), error.size()); + EXPECT_THAT(model, NotNull()) << error.data(); + EXPECT_FLOAT_EQ(model->mesh_texcoord[2*model->mesh_facetexcoord[ 0]], .0); + EXPECT_FLOAT_EQ(model->mesh_texcoord[2*model->mesh_facetexcoord[ 1]], .2); + EXPECT_FLOAT_EQ(model->mesh_texcoord[2*model->mesh_facetexcoord[ 2]], .1); + EXPECT_FLOAT_EQ(model->mesh_texcoord[2*model->mesh_facetexcoord[ 3]], .0); + EXPECT_FLOAT_EQ(model->mesh_texcoord[2*model->mesh_facetexcoord[ 4]], .1); + EXPECT_FLOAT_EQ(model->mesh_texcoord[2*model->mesh_facetexcoord[ 5]], .3); + EXPECT_FLOAT_EQ(model->mesh_texcoord[2*model->mesh_facetexcoord[ 6]], .2); + EXPECT_FLOAT_EQ(model->mesh_texcoord[2*model->mesh_facetexcoord[ 7]], .0); + EXPECT_FLOAT_EQ(model->mesh_texcoord[2*model->mesh_facetexcoord[ 8]], .3); + EXPECT_FLOAT_EQ(model->mesh_texcoord[2*model->mesh_facetexcoord[ 9]], .1); + EXPECT_FLOAT_EQ(model->mesh_texcoord[2*model->mesh_facetexcoord[10]], .2); + EXPECT_FLOAT_EQ(model->mesh_texcoord[2*model->mesh_facetexcoord[11]], .3); + mj_deleteModel(model); +} + +TEST_F(MjCMeshTest, MissingTexCoord) { + static constexpr char xml[] = R"( + + + + + + )"; + std::array error; + mjModel* model = LoadModelFromString(xml, error.data(), error.size()); + EXPECT_THAT(model, testing::IsNull()); + EXPECT_THAT(error.data(), HasSubstr("texcoord must be 2*nv")); +} + + } // namespace } // namespace mujoco From fb8f77df74cfa661a3deb4cf9c2d5113362ca446 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Fri, 26 Jan 2024 06:12:33 -0800 Subject: [PATCH 33/92] Throw error if qhull is called with NaNs. PiperOrigin-RevId: 601745538 Change-Id: I590673ca7975e4abc171f98704a6c8f8d08f8327 --- src/user/user_mesh.cc | 4 ++++ test/user/user_mesh_test.cc | 22 ++++++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/src/user/user_mesh.cc b/src/user/user_mesh.cc index 827f8174..5da757ba 100644 --- a/src/user/user_mesh.cc +++ b/src/user/user_mesh.cc @@ -1506,6 +1506,10 @@ void mjCMesh::MakeGraph(void) { throw mjCError(this, "could not allocate data for qhull"); } for (int i=0; i<3*nvert_; i++) { + if (!std::isfinite(vert_[i])) { + mju_free(data); + throw mjCError(this, "vertex coordinate %d is not finite", NULL, i); + } data[i] = (double)vert_[i]; } diff --git a/test/user/user_mesh_test.cc b/test/user/user_mesh_test.cc index 497d25d0..a4bd0052 100644 --- a/test/user/user_mesh_test.cc +++ b/test/user/user_mesh_test.cc @@ -923,6 +923,28 @@ TEST_F(MjCMeshTest, MissingTexCoord) { EXPECT_THAT(error.data(), HasSubstr("texcoord must be 2*nv")); } +// ----------------------------- qhull ---------------------------------------- + +TEST_F(MjCMeshTest, NaNConvexHullDisallowed) { + static constexpr char xml[] = R"( + + + + + + )"; + static char warning[1024]; + warning[0] = '\0'; + mju_user_warning = [](const char* msg) { + util::strcpy_arr(warning, msg); + }; + std::array error; + mjModel* model = LoadModelFromString(xml, error.data(), error.size()); + EXPECT_THAT(model, testing::IsNull()); + EXPECT_THAT(error.data(), HasSubstr("vertex coordinate 0 is not finite")); + mj_deleteModel(model); +} + } // namespace } // namespace mujoco From ee071242c63339e33943f3bcd2d6e2f45a84d1b2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?B=C3=A1lint=20Hodossy?= Date: Fri, 26 Jan 2024 14:35:03 +0000 Subject: [PATCH 34/92] Constrain MjScene EventHandlers as events, consistently use MjStepArgs --- unity/Runtime/Components/MjScene.cs | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/unity/Runtime/Components/MjScene.cs b/unity/Runtime/Components/MjScene.cs index 6d80893e..c55c5cf5 100644 --- a/unity/Runtime/Components/MjScene.cs +++ b/unity/Runtime/Components/MjScene.cs @@ -78,11 +78,11 @@ public class MjScene : MonoBehaviour { private List _orderedComponents; - public EventHandler postInitEvent; - public EventHandler preUpdateEvent; - public EventHandler ctrlCallback; - public EventHandler postUpdateEvent; - public EventHandler preDestroyEvent; + public event EventHandler postInitEvent; + public event EventHandler preUpdateEvent; + public event EventHandler ctrlCallback; + public event EventHandler postUpdateEvent; + public event EventHandler preDestroyEvent; protected unsafe void Start() { SceneRecreationAtLateUpdateRequested = false; @@ -94,9 +94,9 @@ public class MjScene : MonoBehaviour { } protected unsafe void FixedUpdate() { - preUpdateEvent?.Invoke(this, EventArgs.Empty); + preUpdateEvent?.Invoke(this, new MjStepArgs(Model, Data)); StepScene(); - postUpdateEvent?.Invoke(this, EventArgs.Empty); + postUpdateEvent?.Invoke(this, new MjStepArgs(Model, Data)); } public bool SceneRecreationAtLateUpdateRequested = false; From 669ee0ee4845b235852f247a1a6806c979129142 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Fri, 26 Jan 2024 07:02:00 -0800 Subject: [PATCH 35/92] Fix migration instructions for `option/collision`. Fixes #1367. PiperOrigin-RevId: 601754379 Change-Id: Ifb0b8fab4da44e161e463a71eec050c6635f3340 --- doc/changelog.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 40b74e01..2b1befd2 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -293,7 +293,7 @@ General - For models which have ``