Files
Motrixlab/scripts/terrain_editor.py

431 lines
20 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/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()