chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #1268 from Balint-H:feature/unity-hmap-dyn

Browse Source 
PiperOrigin-RevId: 598652839
Change-Id: I3285f041c5c2e3e9286f1bb106825ee20e7147e3
This commit is contained in:
Copybara-Service
2024-01-15 12:54:06 -08:00
parent 2388b651e8 2c024164ea
commit 3b440921df
2 changed files with 145 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
unity/Runtime/Components/Shapes/MjHeightFieldShape.cs
+121 -25
View File
@@ -17,52 +17,148 @@ using System.IO;
using System.Linq;
using System.Xml;
using UnityEngine;
namespace Mujoco
{
// internal imports
namespace Mujoco {
[Serializable]
public class MjHeightFieldShape : IMjShape
{
public Terrain terrain;
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.")]
public string heightMapExportPath;
[Tooltip("The path, relative to Application.dataPath, where the heightmap " +
"data will be saved/exported in PNG format. Leave blank if hfield data should " +
"be set instead programmatically (faster).")]
public string HeightMapExportPath;
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 Vector3 HeightMapScale => terrain.terrainData.heightmapScale;
public bool ExportImage => !string.IsNullOrEmpty(HeightMapExportPath);
public void ToMjcf(XmlElement mjcf, Transform transform){
ExportHeightMap();
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;
[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. " +
"If nonzero, increasing this can improve performance.")]
public int UpdateLimit;
private int _updateCountdown;
[HideInInspector]
public float MinimumHeight { get; private set; }
[HideInInspector]
public float MaximumHeight { get; private set; }
public int HeightFieldId { get; private set; }
public unsafe void ToMjcf(XmlElement mjcf, Transform transform) {
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 - 1) * HeightMapScale.x / 2f,
Terrain.transform.localPosition.y,
-(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 - 1) * HeightMapScale.x / 2,
Terrain.transform.localPosition.y,
-(HeightMapWidth - 1) * HeightMapScale.z / 2);
}
var scene = MjScene.Instance;
var assetName = scene.GenerationContext.AddHeightFieldAsset(this);
if (Application.isPlaying) {
scene.postInitEvent += (unused_first, unused_second) =>
HeightFieldId =
MujocoLib.mj_name2id(scene.Model, (int)MujocoLib.mjtObj.mjOBJ_HFIELD, assetName);
}
if (UpdateLimit > 0) {
_updateCountdown = UpdateLimit;
if (UpdateLimit > 1) {
scene.preUpdateEvent += (unused_first, unused_second) => CountdownUpdateCondition();
}
TerrainCallbacks.heightmapChanged += (unused_first, unused_second, unused_third) =>
RebuildHeightField();
}
mjcf.SetAttribute("hfield", assetName);
PrepareHeightMap();
}
public void FromMjcf(XmlElement mjcf){
public void FromMjcf(XmlElement mjcf) {
}
public void ExportHeightMap(){
RenderTexture.active = terrain.terrainData.heightmapTexture;
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), 0, 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;
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 += (unused_first, unused_second) => 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;
public Vector4 GetChangeStamp(){
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() {
if (_updateCountdown < 1) return;
_updateCountdown -= 1;
}
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;
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() {
return Vector4.one;
}
public Tuple<Vector3[], int[]> BuildMesh(){
return Tuple.Create(new Vector3[]{}, new int[]{});
public Tuple<Vector3[], int[]> BuildMesh() {
return null;
}
public void DebugDraw(Transform transform){
}
public void DebugDraw(Transform transform) {}
}
}
unity/Runtime/Tools/MjcfGenerationContext.cs
+24 -10
View File
@@ -34,7 +34,9 @@ public class MjcfGenerationContext {
private int _nuserSensor;
private int _numGeneratedNames = 0;
private Dictionary<Mesh, string> _meshAssets = new Dictionary<Mesh, string>();
private Dictionary<MjHeightFieldShape, string> _hFieldAssets = new Dictionary<MjHeightFieldShape, string>();
private Dictionary<MjHeightFieldShape, string> _hFieldAssets =
new Dictionary<MjHeightFieldShape, string>();
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);
@@ -116,17 +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);
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}"));
$@"{hFieldComponent.HeightMapScale.x * (hFieldComponent.HeightMapLength - 1) / 2f} {
hFieldComponent.HeightMapScale.z * (hFieldComponent.HeightMapWidth - 1) / 2f} {
heightRange} {
baseHeight}"));
}
}
}