Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[FEATURE] Support Equality Constraint and Closed-loop robots #636

Merged
merged 20 commits into from
Feb 18, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions examples/rigid/closed_loop.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
import genesis as gs

gs.init(backend=gs.cpu)

scene = gs.Scene(
viewer_options=gs.options.ViewerOptions(
camera_pos=(10, 0, 10),
camera_lookat=(0.0, 0.0, 3),
camera_fov=60,
),
)

franka = scene.add_entity(
gs.morphs.MJCF(
file="xml/four_bar_linkage.xml",
),
)

scene.build()
for i in range(10000):
scene.step()
1 change: 1 addition & 0 deletions genesis/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -247,6 +247,7 @@ def _gs_exit():
IntEnum,
JOINT_TYPE,
GEOM_TYPE,
EQUALITY_TYPE,
CTRL_MODE,
PARA_LEVEL,
ACTIVE,
Expand Down
50 changes: 50 additions & 0 deletions genesis/assets/xml/four_bar_linkage.xml
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
<mujoco>
<option gravity="0 0 0" timestep="0.002" />
<asset>
<material name="blue_" rgba="0 0 1 1" />
<material name="green" rgba="0 1 0 1" />
<material name="red__" rgba="1 0 0 1" />
<material name="white" rgba="1 1 1 1" />
</asset>

<default>
<default class="visual">
<geom type="mesh" contype="0" conaffinity="0" group="2"/>
</default>
</default>

<worldbody>
<geom type="plane" size="10 10 0.1" pos="0 0 -2" rgba=".9 0 0 1" />
<light diffuse=".5 .5 .5" pos="0 0 3" dir="0 0 -1" />
<body name="link_1" pos="0 0 0">
<geom type="cylinder" size=".2 2" pos="0 0 2" euler="0 0 0" material="red__" class="visual"/>
<geom type="cylinder" size=".25 .25" pos="0 0 4" euler="0 90 0" material="red__" class="visual"/>
<geom type="cylinder" size=".25 .25" pos="0 0 0" euler="0 90 0" material="red__" class="visual"/>
<body name="link_2" pos="0.5 0 0" euler="0 0 0">
<joint name="hinge_1" pos="0 0 0" axis="1 0 0" />
<geom type="cylinder" size=".2 2" pos="0 2 0" euler="90 0 0" material="blue_" class="visual"/>
<geom type="cylinder" size=".25 .25" pos="0 4 0" euler="0 90 0" material="blue_" class="visual"/>
<geom type="cylinder" size=".25 .25" pos="0 0 0" euler="0 90 0" material="blue_" class="visual"/>
<body name="link_3" pos="-0.5 4 0" euler="0 0 0">
<joint name="hinge_2" pos="0 0 0" axis="1 0 0" />
<geom type="cylinder" size=".2 2" pos="0 0 2" euler="0 0 0" material="green" class="visual"/>
<geom type="cylinder" size=".25 .25" pos="0 0 0" euler="0 90 0" material="green" class="visual"/>
<geom type="cylinder" size=".25 .25" pos="0 0 4" euler="0 90 0" material="green" class="visual"/>
<body name="link_4" pos="0.5 0 4" euler="0 0 0">
<joint name="hinge_3" pos="0 0 0" axis="1 0 0" />
<geom type="cylinder" size=".2 2" pos="0 -2 0" euler="90 0 0" material="white" class="visual"/>
<geom type="cylinder" size=".25 .25" pos="0 0 0" euler="0 90 0" material="white" class="visual"/>
<geom type="cylinder" size=".25 .25" pos="0 -4 0" euler="0 90 0" material="white" class="visual"/>
</body>
</body>
</body>
</body>
</worldbody>
<equality>
<connect name="kinematic_link1" active="true" body1="link_1" body2="link_4" anchor="0 0 4" />
</equality>
<actuator>
<motor joint="hinge_1" ctrlrange="-100000 100000" ctrllimited="true" />
<motor joint="hinge_2" ctrlrange="-100000 100000" ctrllimited="true" />
</actuator>
</mujoco>
6 changes: 6 additions & 0 deletions genesis/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,12 @@ class JOINT_TYPE(IntEnum):
PLANAR = 5


class EQUALITY_TYPE(IntEnum):
CONNECT = 0
WELD = 1
JOINT = 2


# joint type in rigid solver, ranked by number of dofs
class CTRL_MODE(IntEnum):
FORCE = 0
Expand Down
68 changes: 67 additions & 1 deletion genesis/engine/entities/rigid_entity/rigid_entity.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
from ..base_entity import Entity
from .rigid_joint import RigidJoint
from .rigid_link import RigidLink
from .rigid_equality import RigidEquality


@ti.data_oriented
Expand Down Expand Up @@ -44,6 +45,7 @@ def __init__(
vgeom_start=0,
vvert_start=0,
vface_start=0,
equality_start=0,
visualize_contact=False,
):
super().__init__(idx, scene, morph, solver, material, surface)
Expand All @@ -62,6 +64,7 @@ def __init__(
self._vgeom_start = vgeom_start
self._vvert_start = vvert_start
self._vface_start = vface_start
self._equality_start = equality_start

self._visualize_contact = visualize_contact

Expand All @@ -74,6 +77,7 @@ def __init__(
def _load_model(self):
self._links = gs.List()
self._joints = gs.List()
self._equalities = gs.List()

if isinstance(self._morph, gs.morphs.Mesh):
self._load_mesh(self._morph, self._surface)
Expand Down Expand Up @@ -362,7 +366,7 @@ def _load_MJCF(self, morph, surface):
q_offset += j_info["n_qs"]
dof_offset += j_info["n_dofs"]

l_infos, j_infos, links_g_info = uu._order_links(l_infos, j_infos, links_g_info)
l_infos, j_infos, links_g_info, ordered_links_idx = uu._order_links(l_infos, j_infos, links_g_info)
for i_l in range(len(l_infos)):
l_info = l_infos[i_l]
j_info = j_infos[i_l]
Expand Down Expand Up @@ -391,6 +395,21 @@ def _load_MJCF(self, morph, surface):
gs.logger.warning("Overriding base link's quat with user provided value in morph.")
self._add_by_info(l_world_info, j_world_info, world_g_info, morph, surface)

for i_e in range(mj.neq):
e_info = mju.parse_equality(mj, i_e, morph.scale, ordered_links_idx)
if e_info["type"] == gs.EQUALITY_TYPE.CONNECT: # only this type is supported right now
self._add_equality(
name=e_info["name"],
type=e_info["type"],
link1_idx=e_info["link1_idx"],
link2_idx=e_info["link2_idx"],
anchor1_pos=e_info["anchor1_pos"],
anchor2_pos=e_info["anchor2_pos"],
rel_pose=e_info["rel_pose"],
torque_scale=e_info["torque_scale"],
sol_params=e_info["sol_params"],
)

def _load_URDF(self, morph, surface):
l_infos, j_infos = uu.parse_urdf(morph, surface)

Expand Down Expand Up @@ -590,6 +609,14 @@ def _add_joint(
dofs_force_range,
init_qpos,
):
if (
len(np.array(dofs_sol_params).shape) == 2
and np.array(dofs_sol_params).shape[0] == 1
and (dofs_sol_params[0][3] >= 1.0 or dofs_sol_params[0][2] >= dofs_sol_params[0][3])
):
gs.logger.warning(f"Joint {name}'s sol_params {dofs_sol_params[0]} look not right, change to default.")
dofs_sol_params = gu.default_solver_params(1)

joint = RigidJoint(
entity=self,
name=name,
Expand Down Expand Up @@ -617,6 +644,25 @@ def _add_joint(
self._joints.append(joint)
return joint

def _add_equality(
self, name, type, link1_idx, link2_idx, anchor1_pos, anchor2_pos, rel_pose, torque_scale, sol_params
):
equality = RigidEquality(
entity=self,
name=name,
idx=self.n_equalities + self._equality_start,
type=type,
link1_idx=link1_idx + self._link_start,
link2_idx=link2_idx + self._link_start,
anchor1_pos=anchor1_pos,
anchor2_pos=anchor2_pos,
rel_pose=rel_pose,
torque_scale=torque_scale,
sol_params=sol_params,
)
self._equalities.append(equality)
return equality

# ------------------------------------------------------------------------------------
# --------------------------------- Jacobian & IK ------------------------------------
# ------------------------------------------------------------------------------------
Expand Down Expand Up @@ -2665,6 +2711,26 @@ def base_joint(self):
"""The base joint of the entity"""
return self._joints[0]

@property
def n_equalities(self):
"""The number of equality constraints in the entity."""
return len(self._equalities)

@property
def equality_start(self):
"""The index of the entity's first RigidEquality in the scene."""
return self._equality_start

@property
def equality_end(self):
"""The index of the entity's last RigidEquality in the scene *plus one*."""
return self._equality_start + self.n_equalities

@property
def equalities(self):
"""The list of equality constraints (`RigidEquality`) in the entity."""
return self._equalities

@property
def is_free(self):
"""Whether the entity is free to move."""
Expand Down
145 changes: 145 additions & 0 deletions genesis/engine/entities/rigid_entity/rigid_equality.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,145 @@
import taichi as ti
import torch

import genesis as gs
import genesis.utils.geom as gu
from genesis.repr_base import RBC


@ti.data_oriented
class RigidEquality(RBC):
"""
Equality class for rigid body entities.
"""

def __init__(
self,
entity,
name,
idx,
type,
link1_idx,
link2_idx,
anchor1_pos,
anchor2_pos,
rel_pose,
torque_scale,
sol_params,
):
self._name = name
self._entity = entity
self._solver = entity.solver

self._uid = gs.UID()
self._idx = idx
self._type = type

self._link1_idx = link1_idx
self._link2_idx = link2_idx
self._anchor1_pos = anchor1_pos
self._anchor2_pos = anchor2_pos
self._rel_pose = rel_pose
self._torque_scale = torque_scale
self._sol_params = sol_params

# ------------------------------------------------------------------------------------
# ----------------------------------- properties -------------------------------------
# ------------------------------------------------------------------------------------

@property
def uid(self):
"""
Returns the unique id of the equality.
"""
return self._uid

@property
def name(self):
"""
Returns the name of the equality.
"""
return self._name

@property
def entity(self):
"""
Returns the entity that the equality belongs to.
"""
return self._entity

@property
def solver(self):
"""
The RigidSolver object that the equality belongs to.
"""
return self._solver

@property
def idx(self):
"""
Returns the global index of the equality in the rigid solver.
"""
return self._idx

@property
def idx_local(self):
"""
Returns the local index of the equality in the entity.
"""
return self._idx - self._entity._equality_start

@property
def type(self):
"""
Returns the type of the equality.
"""
return self._type

@property
def link1_idx(self):
"""
Returns the index of the first link.
"""
return self._link1_idx

@property
def link2_idx(self):
"""
Returns the index of the second link.
"""
return self._link2_idx

@property
def anchor1_pos(self):
"""
Returns the position of the first anchor.
"""
return self._anchor1_pos

@property
def anchor2_pos(self):
"""
Returns the position of the second anchor.
"""
return self._anchor2_pos

@property
def rel_pose(self):
"""
Returns the relative pose between the two links.
"""
return self._rel_pose

@property
def torque_scale(self):
"""
Returns the torque scale of the equality.
"""
return self._torque_scale

@property
def sol_params(self):
"""
Returns the solver parameters of the equality.
"""
return self._sol_params
Loading