Skip to content

Commit

Permalink
Merge pull request #1268 from Balint-H:feature/unity-hmap-dyn
Browse files Browse the repository at this point in the history 
PiperOrigin-RevId: 598652839
Change-Id: I3285f041c5c2e3e9286f1bb106825ee20e7147e3
  • Loading branch information
@copybara-github
copybara-github committed Jan 15, 2024
2 parents 2388b65 + 2c02416 commit 3b44092
Show file tree
Hide file tree
Showing 2 changed files with 144 additions and 34 deletions.
144 changes: 120 additions & 24 deletions unity/Runtime/Components/Shapes/MjHeightFieldShape.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,52 +17,148 @@
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 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));

public int HeightMapWidth => Terrain.terrainData.heightmapTexture.width;
public int HeightMapLength => Terrain.terrainData.heightmapTexture.height;
public Vector3 HeightMapScale => Terrain.terrainData.heightmapScale;

[Tooltip("The path, relative to Application.dataPath, where the heightmap data will be save in PNG format.")]
public string 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. " +
"If nonzero, increasing this can improve performance.")]
public int UpdateLimit;

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;
private int _updateCountdown;

public void ToMjcf(XmlElement mjcf, Transform transform){
ExportHeightMap();
[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 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);
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;
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 void ExportHeightMap(){
RenderTexture.active = terrain.terrainData.heightmapTexture;
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);
texture.ReadPixels(new Rect(0, 0, RenderTexture.active.width, RenderTexture.active.height),
0,
0);
RenderTexture.active = null;
File.WriteAllBytes(FullHeightMapPath, texture.EncodeToPNG());

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 Vector4 GetChangeStamp(){
return Vector4.one;
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 Tuple<Vector3[], int[]> BuildMesh(){
return Tuple.Create(new Vector3[]{}, new int[]{});
public Vector4 GetChangeStamp() {
return Vector4.one;
}

public void DebugDraw(Transform transform){
public Tuple<Vector3[], int[]> BuildMesh() {
return null;
}

public void DebugDraw(Transform transform) {}
}
}
34 changes: 24 additions & 10 deletions unity/Runtime/Tools/MjcfGenerationContext.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,9 @@ public int NUserSensor {
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);
Expand Down Expand Up @@ -98,7 +100,7 @@ private void GenerateAssetsMjcf(XmlElement mjcf) {
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);
Expand All @@ -116,17 +118,29 @@ private static void GenerateMeshMjcf(Mesh mesh, 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);
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}"));
}
}
}

0 comments on commit 3b44092

Please sign in to comment.