From e566e67cd3d11f565ac70757713f105dac7f096f Mon Sep 17 00:00:00 2001
From: Sylvester Kaczmarek
 <16242628+sylvesterkaczmarek@users.noreply.github.com>
Date: Sat, 22 Aug 2026 10:36:31 +0100
Subject: [PATCH] Normalize Unity heightfield data for MuJoCo

---
 .../Components/Shapes/MjHeightFieldShape.cs   | 22 ++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs b/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs
index 4ba8de83261..0eb05bf0ea1 100644
--- a/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs
+++ b/unity/Runtime/Components/Shapes/MjHeightFieldShape.cs
@@ -45,6 +45,8 @@ public class MjHeightFieldShape : IMjShape {
   public int UpdateLimit;
 
   private int _updateCountdown;
+  private float _minimumHeightMapValue;
+  private float _maximumHeightMapValue;
 
   [HideInInspector]
   public float MinimumHeight { get; private set; }
@@ -102,12 +104,15 @@ public void PrepareHeightMap() {
     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;
+    var pixels = texture.GetPixels();
+    _maximumHeightMapValue = pixels.Max(c => c.r);
+    _minimumHeightMapValue = pixels.Min(c => c.r);
+    MaximumHeight = _maximumHeightMapValue * HeightMapScale.y * 2;
+    MinimumHeight = _minimumHeightMapValue * HeightMapScale.y * 2;
 
     RenderTexture.active = null;
     if (ExportImage) {
-      if (minimumHeight > 0.0001)
+      if (_minimumHeightMapValue > 0.0001)
         Debug.LogWarning("Due to assumptions in MuJoCo heightfields, terrains should have a " +
                          "minimum heightmap value of 0.");
       File.WriteAllBytes(FullHeightMapPath, texture.EncodeToPNG());
@@ -124,8 +129,15 @@ public unsafe void UpdateHeightFieldData() {
         0);
     RenderTexture.active = null;
 
-    float[] curData = texture.GetPixels(0, 0, texture.width, texture.height)
-        .Select(c => c.r * 2).ToArray();
+    var pixels = texture.GetPixels(0, 0, texture.width, texture.height);
+    float heightRange = _maximumHeightMapValue - _minimumHeightMapValue;
+    float[] curData;
+    if (heightRange > Mathf.Epsilon) {
+      curData = pixels.Select(
+          c => Mathf.Clamp01((c.r - _minimumHeightMapValue) / heightRange)).ToArray();
+    } else {
+      curData = new float[pixels.Length];
+    }
     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];