fix: clamp std before distribution, lower init_noise to 0.5, NaN guard

scripts/terrain_editor.py

@@ -0,0 +1,430 @@
+#!/usr/bin/env python3
+"""Terrain editor GUI — draw pyramids, mark spawn zones, export PNG + coordinates.
+
+Usage:
+    uv run scripts/terrain_editor.py
+"""
+import tkinter as tk
+from tkinter import ttk, messagebox
+import numpy as np, os, cv2
+
+# ═══ defaults ═══
+HS = 0.05; VS = 0.005
+CELL_M = 8.0; BORDER_M = 5.0
+PLATFORM_M = 1.0            # platform 1m
+SPAWN_RADIUS_M = 0.5        # spawn zone ±0.5m around center
+DEFAULT_STEP_H = 0.20; DEFAULT_STEP_D = 0.20; DEFAULT_NUM_STEPS = 10
+DEFAULT_REF_PLANE_CM = 200  # all cells start from same reference height
+CELL_PX = int(CELL_M / HS); BORDER_PX = int(BORDER_M / HS)
+PLATFORM_PX = int(PLATFORM_M / HS); SPAWN_RADIUS_PX = int(SPAWN_RADIUS_M / HS)
+
+
+class TerrainEditor:
+    def __init__(self, root):
+        self.root = root
+        self.root.title("Terrain Editor")
+        self.rows = 2; self.cols = 4
+        self.cell_types = {}
+        self._init_defaults()
+        self.selected = (0, 0)
+        self._dragging = False
+        self._build_ui()
+        self._sync_params()
+        self._redraw_all()
+
+    def _init_defaults(self):
+        for r in range(self.rows):
+            for c in range(self.cols):
+                if r == 0:
+                    self.cell_types[(r, c)] = {"type": "flat", "spawn": True, "level": 0,
+                                                "ref_plane_cm": DEFAULT_REF_PLANE_CM}
+                else:
+                    self.cell_types[(r, c)] = {
+                        "type": "convex" if c % 2 == 0 else "concave",
+                        "step_h": DEFAULT_STEP_H, "step_d": DEFAULT_STEP_D,
+                        "num_steps": DEFAULT_NUM_STEPS, "spawn": True, "level": 1,
+                        "ref_plane_cm": DEFAULT_REF_PLANE_CM,
+                    }
+
+    # ═══ UI ═══
+    def _build_ui(self):
+        paned = ttk.PanedWindow(self.root, orient=tk.HORIZONTAL)
+        paned.pack(fill=tk.BOTH, expand=True)
+        left = ttk.Frame(paned); paned.add(left, weight=2)
+        right = ttk.Frame(paned); paned.add(right, weight=1)
+        self._build_preview_ui(left)
+        self._build_params_ui(right)
+
+    def _build_preview_ui(self, parent):
+        ttk.Label(parent, text="Terrain Preview (click to select cell, right-click toggle spawn)", font=("", 10)).pack(pady=2)
+        self.info_label = ttk.Label(parent, text="")
+        self.info_label.pack()
+        self.preview = tk.Canvas(parent, bg="#333", width=600, height=400)
+        self.preview.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
+        self.preview.bind("<Button-1>", self._on_click)
+        self.preview.bind("<B1-Motion>", self._on_drag)
+        self.preview.bind("<Button-3>", self._on_right_click)
+        ttk.Label(parent, text="Left-click: select | Right-click: toggle spawn | Drag: select").pack()
+        ctrl = ttk.Frame(parent)
+        ctrl.pack(pady=5)
+        ttk.Label(ctrl, text="Rows:").pack(side=tk.LEFT)
+        self.rows_var = tk.IntVar(value=self.rows)
+        ttk.Spinbox(ctrl, from_=1, to=10, width=4, textvariable=self.rows_var,
+                    command=self._on_grid_size).pack(side=tk.LEFT, padx=2)
+        ttk.Label(ctrl, text="Cols:").pack(side=tk.LEFT, padx=(10,0))
+        self.cols_var = tk.IntVar(value=self.cols)
+        ttk.Spinbox(ctrl, from_=1, to=10, width=4, textvariable=self.cols_var,
+                    command=self._on_grid_size).pack(side=tk.LEFT, padx=2)
+        ttk.Button(parent, text="Export PNG + Coords", command=self._export).pack(pady=5)
+
+    def _build_params_ui(self, parent):
+        f = ttk.Frame(parent); f.pack(padx=10, pady=5, fill=tk.X)
+        ttk.Label(f, text="Cell Type:").grid(row=0, column=0, sticky=tk.W)
+        self.type_var = tk.StringVar(value="flat")
+        ttk.Combobox(f, textvariable=self.type_var, values=["flat", "convex", "concave"],
+                     state="readonly", width=10).grid(row=0, column=1, padx=5)
+        self.type_var.trace("w", lambda *a: self._on_param_change())
+        ttk.Label(f, text="Level:").grid(row=0, column=2, sticky=tk.W, padx=(20,0))
+        self.level_var = tk.IntVar(value=0)
+        ttk.Spinbox(f, from_=0, to=9, width=3, textvariable=self.level_var,
+                    command=self._on_param_change).grid(row=0, column=3)
+        self.spawn_var = tk.BooleanVar(value=True)
+        ttk.Checkbutton(f, text="Spawn", variable=self.spawn_var,
+                        command=self._on_param_change).grid(row=0, column=4, padx=10)
+
+        ttk.Separator(parent, orient=tk.HORIZONTAL).pack(fill=tk.X, pady=5, padx=10)
+        ttk.Label(parent, text="Pyramid Params").pack()
+        g = ttk.Frame(parent); g.pack(padx=10, pady=5, fill=tk.X)
+        ttk.Label(g, text="Step height (m):").grid(row=0, column=0, sticky=tk.W)
+        self.sh_var = tk.StringVar(value=str(DEFAULT_STEP_H))
+        ttk.Entry(g, textvariable=self.sh_var, width=7).grid(row=0, column=1, padx=5)
+        self.sh_var.trace("w", lambda *a: self._on_param_change())
+        ttk.Label(g, text="Step tread (m):").grid(row=1, column=0, sticky=tk.W)
+        self.sd_var = tk.StringVar(value=str(DEFAULT_STEP_D))
+        ttk.Entry(g, textvariable=self.sd_var, width=7).grid(row=1, column=1, padx=5)
+        self.sd_var.trace("w", lambda *a: self._on_param_change())
+        ttk.Label(g, text="Num steps:").grid(row=2, column=0, sticky=tk.W)
+        self.ns_var = tk.StringVar(value=str(DEFAULT_NUM_STEPS))
+        ttk.Entry(g, textvariable=self.ns_var, width=7).grid(row=2, column=1, padx=5)
+        ttk.Label(g, text="Ref plane (cm):").grid(row=3, column=0, sticky=tk.W)
+        self.ref_var = tk.StringVar(value=str(DEFAULT_REF_PLANE_CM))
+        ttk.Entry(g, textvariable=self.ref_var, width=7).grid(row=3, column=1, padx=5)
+        self.ns_var.trace("w", lambda *a: self._on_param_change())
+
+        ttk.Separator(parent, orient=tk.HORIZONTAL).pack(fill=tk.X, pady=5, padx=10)
+        ttk.Label(parent, text="Selected Cell").pack()
+        self.cell_label = ttk.Label(parent, text="")
+        self.cell_label.pack()
+
+    # ═══ events ═══
+    def _cell_at(self, ex, ey):
+        m = 5; pw = self.preview.winfo_width(); ph = self.preview.winfo_height()
+        cw = (pw - 2*m) // max(self.cols,1); ch = (ph - 2*m) // max(self.rows,1)
+        col = (ex - m) // cw; row = (ey - m) // ch
+        if 0 <= col < self.cols and 0 <= row < self.rows:
+            return row, col, m + col*cw, m + row*ch, cw, ch
+        return None
+
+    def _on_click(self, event):
+        v = self._cell_at(event.x, event.y)
+        if v:
+            self.selected = (v[0], v[1])
+            self._sync_params(); self._redraw_all()
+
+    def _on_drag(self, event):
+        v = self._cell_at(event.x, event.y)
+        if v:
+            self.selected = (v[0], v[1])
+            self._sync_params(); self._redraw_all()
+
+    def _on_right_click(self, event):
+        v = self._cell_at(event.x, event.y)
+        if v:
+            r, c = v[0], v[1]
+            ct = self.cell_types.get((r, c), {"type": "flat", "spawn": True, "level": r})
+            ct = dict(ct)  # copy before modifying
+            ct["spawn"] = not ct.get("spawn", True)
+            self.cell_types[(r, c)] = ct
+            if (r, c) == self.selected:
+                self._sync_params()
+            self._redraw_all()
+
+    def _on_grid_size(self):
+        try: nr = self.rows_var.get()
+        except: nr = self.rows
+        try: nc = self.cols_var.get()
+        except: nc = self.cols
+        if nr == self.rows and nc == self.cols: return
+        old = self.cell_types
+        self.rows, self.cols = nr, nc
+        self.cell_types = {}
+        for r in range(nr):
+            for c in range(nc):
+                self.cell_types[(r,c)] = old.get((r,c), {"type": "flat", "spawn": True, "level": r})
+        self._redraw_all()
+
+    # ═══ sync ═══
+    def _sync_params(self):
+        ct = self.cell_types.get(self.selected, {"type": "flat", "spawn": True, "level": self.selected[0]})
+        self.type_var.set(ct.get("type", "flat"))
+        self.spawn_var.set(ct.get("spawn", True))
+        r = self.selected[0]
+        self.level_var.set(ct.get("level", r))
+        self.ref_var.set(str(ct.get("ref_plane_cm", DEFAULT_REF_PLANE_CM)))
+        self.sh_var.set(str(ct.get("step_h", DEFAULT_STEP_H)))
+        self.sd_var.set(str(ct.get("step_d", DEFAULT_STEP_D)))
+        self.ns_var.set(str(ct.get("num_steps", DEFAULT_NUM_STEPS)))
+        r, c = self.selected
+        half_x = BORDER_M + self.cols * CELL_M / 2
+        half_y = BORDER_M + self.rows * CELL_M / 2
+        cx = -half_x + BORDER_M + c * CELL_M + CELL_M / 2
+        cy = half_y - BORDER_M - r * CELL_M - CELL_M / 2
+        self.cell_label.config(text=f"({r},{c}) center: x={cx:+.1f} y={cy:+.1f}  type={ct['type']}")
+
+    def _on_param_change(self):
+        r, c = self.selected
+        try: sh = float(self.sh_var.get()); sd = float(self.sd_var.get())
+        except: return
+        try: ns = int(self.ns_var.get())
+        except: return
+        try: ref_cm = float(self.ref_var.get())
+        except: ref_cm = DEFAULT_REF_PLANE_CM
+        cell = {"type": self.type_var.get(), "spawn": self.spawn_var.get(),
+                "level": self.level_var.get(), "ref_plane_cm": ref_cm}
+        if cell["type"] != "flat":
+            cell.update({"step_h": sh, "step_d": sd, "num_steps": ns})
+        self.cell_types[(r, c)] = cell
+        self._redraw_all()
+
+    # ═══ draw ═══
+    def _redraw_all(self):
+        w = self.preview.winfo_width(); h = self.preview.winfo_height()
+        if w < 10: w = 600
+        if h < 10: h = 400
+        self._draw_preview(w, h)
+        half_x = BORDER_M + self.cols * CELL_M / 2
+        half_y = BORDER_M + self.rows * CELL_M / 2
+        self.info_label.config(
+            text=f"{self.rows}×{self.cols}  "
+                 f"{self.cols*CELL_M+2*BORDER_M:.0f}×{self.rows*CELL_M+2*BORDER_M:.0f}m  "
+                 f"spawn_cy = {half_y-BORDER_M-CELL_M/2:.0f} - row*{CELL_M:.0f}")
+
+    def _draw_preview(self, pw, ph):
+        cv = self.preview; cv.delete("all")
+        m = 5; cw = (pw - 2*m) // max(self.cols, 1); ch = (ph - 2*m) // max(self.rows, 1)
+        cw = max(cw, 30); ch = max(ch, 30)
+        colors = {"flat": "#5b8c5a", "convex": "#c0392b", "concave": "#2471a3"}
+
+        for r in range(self.rows):
+            for c in range(self.cols):
+                x1, y1 = m + c*cw, m + r*ch
+                x2, y2 = x1 + cw, y1 + ch
+                ct = self.cell_types.get((r,c), {"type": "flat", "spawn": True, "level": r})
+                cv.create_rectangle(x1, y1, x2, y2, fill=colors.get(ct["type"], "#555"),
+                                    outline="#888", width=1)
+                # Cell center dot
+                cx = (x1+x2)//2; cy = (y1+y2)//2
+                cv.create_oval(cx-3, cy-3, cx+3, cy+3, fill="white", outline="")
+
+                # Pyramid stairs rings
+                if ct["type"] != "flat":
+                    sh = ct.get("step_h", DEFAULT_STEP_H)
+                    sd = ct.get("step_d", DEFAULT_STEP_D)
+                    ns = ct.get("num_steps", DEFAULT_NUM_STEPS)
+                    concave = ct["type"] == "concave"
+                    p2 = max(2, cw // 16)
+                    step_px = max(1, (cw//2 - p2) // max(ns, 1))
+                    h_max = int(sh * ns / VS)
+                    for i in range(ns + 1):
+                        half = p2 + (ns - i) * step_px
+                        if concave:
+                            frac = (ns - i) / max(ns, 1)
+                        else:
+                            frac = i / max(ns, 1)
+                        g = int(180 - frac * 100)
+                        clr = f"#{g:02x}{g:02x}{g:02x}"
+                        cv.create_rectangle(cx - half, cy - half, cx + half, cy + half,
+                                            fill=clr, outline="")
+
+                # Spawn zone (green rect)
+                if ct.get("spawn", True):
+                    sz = max(2, int(SPAWN_RADIUS_M / CELL_M * cw))
+                    cv.create_rectangle(cx - sz, cy - sz, cx + sz, cy + sz,
+                                        outline="#00ff00", width=2)
+
+                # Level + type label
+                lvl = ct.get("level", r)
+                lbl = f"L{lvl} {ct['type'][:3]}"
+                if ct["type"] != "flat":
+                    lbl = f"L{lvl} {ct['type'][:3]}-{sh*100:.0f}cm"
+                cv.create_text(x1 + 20, y1 + 10, text=lbl, fill="white",
+                               font=("", 8), anchor=tk.NW)
+
+        # Highlight selected cell
+        r, c = self.selected
+        x1, y1 = m + c*cw, m + r*ch
+        x2, y2 = x1 + cw, y1 + ch
+        cv.create_rectangle(x1, y1, x2, y2, outline="yellow", width=3)
+
+        # Level labels on right
+        for r in range(self.rows):
+            y = m + r*ch + ch//2
+            cv.create_text(pw - 15, y, text=f"L{r}", fill="white", font=("", 12, "bold"))
+
+    # ═══ generate + export ═══
+    def _generate_png(self):
+        tot_rows = self.rows * CELL_PX + 2 * BORDER_PX
+        tot_cols = self.cols * CELL_PX + 2 * BORDER_PX
+        canvas = np.zeros((tot_rows, tot_cols), dtype=np.uint16)
+        for r in range(self.rows):
+            for c in range(self.cols):
+                x0 = BORDER_PX + c * CELL_PX; y0 = BORDER_PX + r * CELL_PX
+                ct = self.cell_types.get((r,c), {"type": "flat", "level": r})
+                if ct["type"] == "flat": continue
+                sh = ct.get("step_h", DEFAULT_STEP_H)
+                sd = ct.get("step_d", DEFAULT_STEP_D)
+                ns = ct.get("num_steps", DEFAULT_NUM_STEPS)
+                concave = ct["type"] == "concave"
+                ref_cm = ct.get("ref_plane_cm", DEFAULT_REF_PLANE_CM)
+                ref_vs = int(ref_cm / 100.0 / VS)
+                h_vs = int(sh / VS); d_px = int(sd / HS)
+                p2 = PLATFORM_PX // 2
+                cx = x0 + CELL_PX // 2; cy = y0 + CELL_PX // 2
+                cv2.rectangle(canvas, (x0, y0), (x0+CELL_PX, y0+CELL_PX), int(ref_vs), -1)
+                for i in range(ns + 1):
+                    half = p2 + (ns - i) * d_px
+                    x1, y1 = cx - half, cy - half; x2, y2 = cx + half, cy + half
+                    if concave:
+                        h = ref_vs - h_vs * i
+                    else:
+                        h = ref_vs + h_vs * i
+                    cv2.rectangle(canvas, (x1, y1), (x2, y2), int(h), -1)
+        hf_m = canvas.astype(np.float32) * VS
+        z_min, z_max = float(hf_m.min()), float(hf_m.max())
+        z_range = max(z_max - z_min, 0.001)
+        png = ((hf_m - z_min) / z_range * 65535).astype(np.uint16)
+        return png, z_range, z_min
+
+    def _export(self):
+        png, z_range, z_min = self._generate_png()
+        out_dir = os.path.join(os.path.dirname(__file__), "..",
+                               "motrix_envs", "src", "motrix_envs", "locomotion",
+                               "go1", "xmls", "assets")
+        os.makedirs(out_dir, exist_ok=True)
+        out_path = os.path.join(out_dir, "flat_stairs.png")
+        cv2.imwrite(out_path, png)
+        w_m = (self.cols * CELL_PX + 2 * BORDER_PX) * HS
+        h_m = (self.rows * CELL_PX + 2 * BORDER_PX) * HS
+        half_x = BORDER_M + self.cols * CELL_M / 2
+        half_y = BORDER_M + self.rows * CELL_M / 2
+
+        lines = [
+            f"# Terrain: {self.rows}×{self.cols}  {w_m:.0f}×{h_m:.0f}m",
+            f"XML: size=\"{w_m/2:.1f} {h_m/2:.1f} {z_range:.3f} {max(z_min,0.001):.3f}\"",
+            f"dreamwaq.py: terrain_rows={self.rows}  terrain_cols={self.cols}",
+            f"",
+            f"# === Cell centers ===",
+        ]
+        for r in range(self.rows):
+            for c in range(self.cols):
+                cx = -half_x + BORDER_M + c * CELL_M + CELL_M / 2
+                cy = half_y - BORDER_M - r * CELL_M - CELL_M / 2
+                ct = self.cell_types.get((r,c), {"type": "flat", "level": r})
+                lines.append(f"  ({r},{c}): x={cx:+.1f} y={cy:+.1f}  {ct['type']}")
+
+        lines.append("")
+        lines.append("# === Level boundaries (robot out of bounds → reset) ===")
+        level_bounds = {}
+        for r in range(self.rows):
+            for c in range(self.cols):
+                lv = self.cell_types.get((r,c), {"level": r})["level"]
+                if lv not in level_bounds:
+                    level_bounds[lv] = {"rmin": r, "rmax": r, "cmin": c, "cmax": c}
+                else:
+                    b = level_bounds[lv]
+                    b["rmin"] = min(b["rmin"], r)
+                    b["rmax"] = max(b["rmax"], r)
+                    b["cmin"] = min(b["cmin"], c)
+                    b["cmax"] = max(b["cmax"], c)
+        for lv in sorted(level_bounds):
+            b = level_bounds[lv]
+            x_min = -half_x + BORDER_M + b["cmin"] * CELL_M
+            x_max = -half_x + BORDER_M + (b["cmax"] + 1) * CELL_M
+            y_min = half_y - BORDER_M - (b["rmax"] + 1) * CELL_M
+            y_max = half_y - BORDER_M - b["rmin"] * CELL_M
+            lines.append(f"  level {lv}: x=[{x_min:+.1f}, {x_max:+.1f}]  "
+                         f"y=[{y_min:+.1f}, {y_max:+.1f}]  "
+                         f"({b['rmax']-b['rmin']+1}×{b['cmax']-b['cmin']+1} cells)")
+
+        lines.append("")
+        lines.append("# === Spawn positions ===")
+        for lv in sorted(level_bounds):
+            b = level_bounds[lv]
+            spawn_cells = [(r,c) for r in range(b["rmin"], b["rmax"]+1)
+                           for c in range(b["cmin"], b["cmax"]+1)
+                           if self.cell_types.get((r,c), {}).get("spawn", True)]
+            if spawn_cells:
+                lines.append(f"  level {lv}: {len(spawn_cells)} spawn cells")
+                for (rr, cc) in spawn_cells:
+                    cx = -half_x + BORDER_M + cc * CELL_M + CELL_M / 2
+                    cy = half_y - BORDER_M - rr * CELL_M - CELL_M / 2
+                    ct = self.cell_types.get((rr,cc), {})
+                    if ct.get("type") == "flat":
+                        z_plat = 0
+                    else:
+                        ref_cm = ct.get("ref_plane_cm", DEFAULT_REF_PLANE_CM)
+                        sh = ct.get("step_h", 0)
+                        ns = ct.get("num_steps", 0)
+                        concave = ct["type"] == "concave"
+                        z_plat = (ref_cm - ns * sh * 100) / 100.0 if concave else (ref_cm + ns * sh * 100) / 100.0
+                    lines.append(f"    ({rr},{cc}) x={cx:+.1f} y={cy:+.1f}  {ct['type']} "
+                                 f"z_plat={z_plat*100:.0f}cm")
+
+        lines.append("")
+        lines.append("# === Pyramid tread details ===")
+        for r in range(self.rows):
+            for c in range(self.cols):
+                ct = self.cell_types.get((r,c), {"type": "flat"})
+                if ct["type"] == "flat":
+                    continue
+                cx = -half_x + BORDER_M + c * CELL_M + CELL_M / 2
+                cy = half_y - BORDER_M - r * CELL_M - CELL_M / 2
+                sh = ct.get("step_h", DEFAULT_STEP_H)
+                sd = ct.get("step_d", DEFAULT_STEP_D)
+                ns = ct.get("num_steps", DEFAULT_NUM_STEPS)
+                concave = ct["type"] == "concave"
+                ref_cm = ct.get("ref_plane_cm", DEFAULT_REF_PLANE_CM)
+                ref_z = ref_cm / 100.0
+                h_vs = int(sh / VS)
+                p2 = PLATFORM_PX // 2
+                d_px = int(sd / HS)
+                plat_z = (ref_cm - ns * sh * 100) / 100.0 if concave else (ref_cm + ns * sh * 100) / 100.0
+                lines.append(f"  ({r},{c}) {ct['type']}  center=({cx:+.1f}, {cy:+.1f})  "
+                             f"ref_plane={ref_z*100:.0f}cm  platform={plat_z*100:.0f}cm  "
+                             f"step_h={sh*100:.0f}cm  tread={sd*100:.0f}cm  steps={ns}")
+                for i in range(ns + 1):
+                    half_m = (p2 + (ns - i) * d_px) * HS
+                    if concave:
+                        z = ref_z - (h_vs * i) * VS
+                    else:
+                        z = ref_z + (h_vs * i) * VS
+                    ring_type = "platform" if i == ns else "ring"
+                    lines.append(f"    {ring_type} {i}: z={z*100:5.0f}cm  "
+                                 f"half={half_m:.2f}m  "
+                                 f"x=[{cx-half_m:+.1f},{cx+half_m:+.1f}] "
+                                 f"y=[{cy-half_m:+.1f},{cy+half_m:+.1f}]")
+
+        lines.append("")
+        lines.append(f"# Training: DREAMWAQ_TERRAIN=flat_stairs "
+                     f"uv run scripts/train_dreamwaq_rsl.py --level N")
+
+        info = "\n".join(lines)
+        print(info)
+        messagebox.showinfo("Exported", f"{out_path}\n\n{info}")
+
+
+if __name__ == "__main__":
+    root = tk.Tk()
+    root.geometry("900x550")
+    TerrainEditor(root)
+    root.mainloop()