feat: initialize project with robot simulation demo

2025-12-20 06:27:49 +08:00 · 2025-11-29 19:47:29 +08:00
parent 8e958c9322
commit 1b8c2f4ee0
5 changed files with 2309 additions and 0 deletions
--- a/src/main.py
+++ b/src/main.py
@@ -0,0 +1,126 @@
+from typing import cast
+
+import genesis as gs
+import numpy as np
+
+########################## init ##########################
+gs.init(backend=gs.gs_backend.gpu)
+
+try:
+    from genesis.engine.entities import RigidEntity
+except ImportError:
+    raise ImportError("genesis.engine.entities.RigidEntity is not available")
+
+########################## create a scene ##########################
+scene = gs.Scene(
+    viewer_options=gs.options.ViewerOptions(
+        camera_pos=(0, -3.5, 2.5),
+        camera_lookat=(0.0, 0.0, 0.5),
+        camera_fov=30,
+        res=(960, 640),
+        max_FPS=60,
+    ),
+    sim_options=gs.options.SimOptions(
+        dt=0.01,
+    ),
+    show_viewer=True,
+)
+
+########################## entities ##########################
+plane = scene.add_entity(
+    gs.morphs.Plane(),
+)
+franka = scene.add_entity(
+    gs.morphs.MJCF(
+        file="xml/franka_emika_panda/panda.xml",
+    ),
+)
+franka = cast(RigidEntity, franka)
+
+
+########################## build ##########################
+scene.build()
+
+jnt_names = [
+    "joint1",
+    "joint2",
+    "joint3",
+    "joint4",
+    "joint5",
+    "joint6",
+    "joint7",
+    "finger_joint1",
+    "finger_joint2",
+]
+dofs_idx = [franka.get_joint(name).dof_idx_local for name in jnt_names]
+
+############ Optional: set control gains ############
+# set positional gains
+franka.set_dofs_kp(
+    kp=np.array([4500, 4500, 3500, 3500, 2000, 2000, 2000, 100, 100]),
+    dofs_idx_local=dofs_idx,
+)
+# set velocity gains
+franka.set_dofs_kv(
+    kv=np.array([450, 450, 350, 350, 200, 200, 200, 10, 10]),
+    dofs_idx_local=dofs_idx,
+)
+# set force range for safety
+franka.set_dofs_force_range(
+    lower=np.array([-87, -87, -87, -87, -12, -12, -12, -100, -100]),
+    upper=np.array([87, 87, 87, 87, 12, 12, 12, 100, 100]),
+    dofs_idx_local=dofs_idx,
+)
+# Hard reset
+for i in range(150):
+    if i < 50:
+        franka.set_dofs_position(np.array([1, 1, 0, 0, 0, 0, 0, 0.04, 0.04]), dofs_idx)
+    elif i < 100:
+        franka.set_dofs_position(
+            np.array([-1, 0.8, 1, -2, 1, 0.5, -0.5, 0.04, 0.04]), dofs_idx
+        )
+    else:
+        franka.set_dofs_position(np.array([0, 0, 0, 0, 0, 0, 0, 0, 0]), dofs_idx)
+
+    scene.step()
+
+# PD control
+for i in range(1250):
+    if i == 0:
+        franka.control_dofs_position(
+            np.array([1, 1, 0, 0, 0, 0, 0, 0.04, 0.04]),
+            dofs_idx,
+        )
+    elif i == 250:
+        franka.control_dofs_position(
+            np.array([-1, 0.8, 1, -2, 1, 0.5, -0.5, 0.04, 0.04]),
+            dofs_idx,
+        )
+    elif i == 500:
+        franka.control_dofs_position(
+            np.array([0, 0, 0, 0, 0, 0, 0, 0, 0]),
+            dofs_idx,
+        )
+    elif i == 750:
+        # control first dof with velocity, and the rest with position
+        franka.control_dofs_position(
+            np.array([0, 0, 0, 0, 0, 0, 0, 0, 0])[1:],
+            dofs_idx[1:],
+        )
+        franka.control_dofs_velocity(
+            np.array([1.0, 0, 0, 0, 0, 0, 0, 0, 0])[:1],
+            dofs_idx[:1],
+        )
+    elif i == 1000:
+        franka.control_dofs_force(
+            np.array([0, 0, 0, 0, 0, 0, 0, 0, 0]),
+            dofs_idx,
+        )
+    # This is the control force computed based on the given control command
+    # If using force control, it's the same as the given control command
+    print("control force:", franka.get_dofs_control_force(dofs_idx))
+
+    # This is the actual force experienced by the dof
+    print("internal force:", franka.get_dofs_force(dofs_idx))
+
+    scene.step()