lerobot/robot-learning-tutorial

1import plotly.graph_objects as go2import numpy as np3import pandas as pd4 5# Scene parameters (same ranges as the Astro integration)6cx, cy = 1.5, 0.5                 # center7a, b = 1.3, 0.45                  # max extent in x/y (ellipse for anisotropy)8 9# Spiral galaxy parameters10num_points = 3000                 # more dots11num_arms = 3                      # number of spiral arms12num_turns = 2.1                   # number of turns per arm13angle_jitter = 0.12               # angular jitter to fan out the arms14pos_noise = 0.015                 # global position noise15 16# Generate points along spiral arms (Archimedean spiral)17t = np.random.rand(num_points) * (2 * np.pi * num_turns)  # progression along the arm18arm_indices = np.random.randint(0, num_arms, size=num_points)19arm_offsets = arm_indices * (2 * np.pi / num_arms)20 21theta = t + arm_offsets + np.random.randn(num_points) * angle_jitter22 23# Normalized radius (0->center, 1->edge). Power <1 to densify the core24r_norm = (t / (2 * np.pi * num_turns)) ** 0.925 26# Radial/lateral noise that slightly increases with radius27noise_x = pos_noise * (0.8 + 0.6 * r_norm) * np.random.randn(num_points)28noise_y = pos_noise * (0.8 + 0.6 * r_norm) * np.random.randn(num_points)29 30# Elliptic projection31x_spiral = cx + a * r_norm * np.cos(theta) + noise_x32y_spiral = cy + b * r_norm * np.sin(theta) + noise_y33 34# Central bulge (additional points very close to the core)35bulge_points = int(0.18 * num_points)36phi_b = 2 * np.pi * np.random.rand(bulge_points)37r_b = (np.random.rand(bulge_points) ** 2.2) * 0.22  # compact bulge38noise_x_b = (pos_noise * 0.6) * np.random.randn(bulge_points)39noise_y_b = (pos_noise * 0.6) * np.random.randn(bulge_points)40x_bulge = cx + a * r_b * np.cos(phi_b) + noise_x_b41y_bulge = cy + b * r_b * np.sin(phi_b) + noise_y_b42 43# Concatenation44x = np.concatenate([x_spiral, x_bulge])45y = np.concatenate([y_spiral, y_bulge])46 47# Central intensity (for sizes/colors). 1 at center, ~0 at edge48z_spiral = 1 - r_norm49z_bulge = 1 - (r_b / max(r_b.max(), 1e-6))  # very bright bulge50z_raw = np.concatenate([z_spiral, z_bulge])51 52# Sizes: keep the 5..10 scale for consistency53sizes = (z_raw + 1) * 554 55# Remove intermediate filtering: keep all placed points, filter at the very end56 57df = pd.DataFrame({58    "x": x,59    "y": y,60    "z": sizes,  # reused for size+color as before61})62 63def get_label(z):64    if z < 0.25:65        return "smol dot"66    if z < 0.5:67        return "ok-ish dot"68    if z < 0.75:69        return "a dot"70    else:71        return "biiig dot"72 73# Labels based on central intensity74df["label"] = pd.Series(z_raw).apply(get_label)75 76# Rendering order: small points first, big ones after (on top)77df = df.sort_values(by="z", ascending=True).reset_index(drop=True)78 79fig = go.Figure()80 81fig.add_trace(go.Scattergl(82    x=df['x'],83    y=df['y'],84    mode='markers',85    marker=dict(86        size=df['z'],87        color=df['z'],88        colorscale=[89            [0, 'rgb(78, 165, 183)'],90            [0.5, 'rgb(206, 192, 250)'],91            [1, 'rgb(232, 137, 171)']92        ],93        opacity=0.9,94    ),95    customdata=df[["label"]],96    hovertemplate="Dot category: %{customdata[0]}",97    hoverlabel=dict(namelength=0),98    showlegend=False99))100 101fig.update_layout(102    autosize=True,103    paper_bgcolor='rgba(0,0,0,0)',104    plot_bgcolor='rgba(0,0,0,0)',105    showlegend=False,106    margin=dict(l=0, r=0, t=0, b=0),107    xaxis=dict(108        showgrid=False,109        zeroline=False,110        showticklabels=False,111        range=[0, 3]112    ),113    yaxis=dict(114        showgrid=False,115        zeroline=False,116        showticklabels=False,117        scaleanchor="x",118        scaleratio=1,119        range=[0, 1]120    )121)122 123# fig.show()124 125fig.write_html(126    "../app/src/content/fragments/banner.html",127    include_plotlyjs=False,128    full_html=False,129    config={130        'displayModeBar': False,131        'responsive': True,132        'scrollZoom': False,133    }134)