feat: enhance scene management and object handling with new features and improvements

- Updated Sphere, Cylinder, Box classes to support random color generation by accepting an optional color parameter and implementing random color assignment if none is provided.
- Improved documentation formatting in the RbsArm class for better readability and consistency.
- Refactored scene management to integrate both Model and Mesh classes, allowing for more flexible object handling. This includes updating object instantiation and adding comments for future color randomization implementation.
- Added `trimesh` as a dependency in the setup configuration to ensure proper functionality for mesh handling.
- Cleaned up node initialization in the runtime launch setup for improved readability.
- Enhanced the default scene configuration by including new ModelData and MeshData objects, with random color generation configured for MeshData.
- Renamed `MeshObjectData` to `ModelData` for clarity and introduced a new `MeshData` class to specifically handle mesh properties.
- Added a `random_color` attribute to `ObjectRandomizerData` and updated `ObjectData` to include a `color` attribute for better object customization.
- Introduced new gripper data classes for better organization and added a `GripperEnum` to categorize different gripper types.
- Implemented a new Mesh class for handling 3D mesh objects in Gazebo simulation, including methods for calculating inertia and generating SDF strings, with error handling for missing mesh files.
This commit is contained in:
Ilya Uraev 2024-10-04 18:25:09 +03:00
parent 31d2afe6e2
commit b42c2e633e
17 changed files with 330 additions and 63 deletions

View file

@ -3,7 +3,8 @@ from .light import LightData
from .objects import (
BoxObjectData,
CylinderObjectData,
MeshObjectData,
MeshData,
ModelData,
ObjectData,
PlaneObjectData,
SphereObjectData,

View file

@ -1,5 +1,7 @@
from dataclasses import dataclass, field
from env_manager.utils.types import Pose
@dataclass
class ObjectRandomizerData:
@ -37,6 +39,7 @@ class ObjectRandomizerData:
random_spawn_position_update_workspace_centre: bool = field(default=False)
random_spawn_volume: tuple[float, float, float] = field(default=(0.5, 0.5, 0.5))
models_rollouts_num: int = field(default=0)
random_color: bool = field(default=False)
@dataclass
@ -67,6 +70,7 @@ class ObjectData:
relative_to: str = field(default="world")
position: tuple[float, float, float] = field(default=(0.0, 0.0, 0.0))
orientation: tuple[float, float, float, float] = field(default=(1.0, 0.0, 0.0, 0.0))
color: tuple[float, float, float, float] | None = field(default=None)
static: bool = field(default=False)
randomize: ObjectRandomizerData = field(default_factory=ObjectRandomizerData)
@ -90,7 +94,6 @@ class PrimitiveObjectData(ObjectData):
collision: bool = field(default=True)
visual: bool = field(default=True)
color: tuple = field(default=(0.8, 0.8, 0.8, 1.0))
mass: float = field(default=0.1)
@ -169,7 +172,7 @@ class BoxObjectData(PrimitiveObjectData):
@dataclass
class MeshObjectData(ObjectData):
class ModelData(ObjectData):
"""
MeshObjectData defines the specific properties for a mesh-based object in the simulation.
@ -180,3 +183,16 @@ class MeshObjectData(ObjectData):
"""
texture: list[float] = field(default_factory=list)
@dataclass
class MeshData(ModelData):
mass: float = field(default_factory=float)
inertia: tuple[float, float, float, float, float, float] = field(
default_factory=tuple
)
cm: Pose = field(default_factory=Pose)
collision: str = field(default_factory=str)
visual: str = field(default_factory=str)
friction: float = field(default_factory=float)
density: float = field(default_factory=float)

View file

@ -1,4 +1,38 @@
from dataclasses import dataclass, field
from enum import Enum
@dataclass
class ToolData:
name: str = field(default_factory=str)
type: str = field(default_factory=str)
@dataclass
class GripperData(ToolData):
pass
@dataclass
class ParallelGripperData(GripperData):
pass
@dataclass
class MultifingerGripperData(GripperData):
pass
@dataclass
class VacuumGripperData(GripperData):
pass
class GripperEnum(Enum):
parallel = ParallelGripperData
mulrifinger = MultifingerGripperData
vacuum = VacuumGripperData
@dataclass

View file

@ -2,13 +2,13 @@ from enum import Enum
from gym_gz.scenario.model_wrapper import ModelWrapper
from .model import Mesh
from .model import Model
from .mesh import Mesh
from .primitives import Box, Cylinder, Plane, Sphere
from .random_object import RandomObject
from .random_primitive import RandomPrimitive
# Enum для типов объектов
class ObjectType(Enum):
BOX = "box"
SPHERE = "sphere"
@ -16,6 +16,7 @@ class ObjectType(Enum):
PLANE = "plane"
RANDOM_PRIMITIVE = "random_primitive"
RANDOM_MESH = "random_mesh"
MODEL = "model"
MESH = "mesh"
OBJECT_MODEL_CLASSES = {
@ -25,7 +26,8 @@ OBJECT_MODEL_CLASSES = {
ObjectType.PLANE: Plane,
ObjectType.RANDOM_PRIMITIVE: RandomPrimitive,
ObjectType.RANDOM_MESH: RandomObject,
ObjectType.MESH: Mesh,
ObjectType.MODEL: Model,
ObjectType.MESH: Mesh
}

View file

@ -0,0 +1,183 @@
import numpy as np
import trimesh
from gym_gz.scenario import model_wrapper
from gym_gz.utils.scenario import get_unique_model_name
from scenario import core as scenario_core
from scenario import gazebo as scenario_gazebo
from trimesh import Trimesh
from env_manager.utils.types import Point
InertiaTensor = tuple[float, float, float, float, float, float]
class Mesh(model_wrapper.ModelWrapper):
""" """
def __init__(
self,
world: scenario_gazebo.World,
name: str = "object",
type: str = "mesh",
relative_to: str = "world",
position: tuple[float, float, float] = (0, 0, 0),
orientation: tuple[float, float, float, float] = (1, 0, 0, 0),
color: tuple[float, float, float, float] | None = None,
static: bool = False,
texture: list[float] = [],
mass: float = 0.0,
inertia: InertiaTensor | None = None,
cm: Point | None = None,
collision: str = "",
visual: str = "",
friction: float = 0.0,
density: float = 0.0,
**kwargs,
):
# Get a unique model name
model_name = get_unique_model_name(world, name)
# Initial pose
initial_pose = scenario_core.Pose(position, orientation)
# Calculate inertia parameters for provided mesh file if not exist
if not inertia and collision and density:
mass, cm, inertia = self.calculate_inertia(collision, density)
else:
raise RuntimeError(
f"Please provide collision mesh filepath for model {name} to calculate inertia"
)
if not color:
color = tuple(np.random.uniform(0.0, 1.0, (3,)))
color = (color[0], color[1], color[2], 1.0)
# Create SDF string for the model
sdf = self.get_sdf(
name=name,
static=static,
collision=collision,
visual=visual,
mass=mass,
inertia=inertia,
friction=friction,
color=color,
center_of_mass=cm,
gui_only=True,
)
# Insert the model
ok_model = world.to_gazebo().insert_model_from_string(
sdf, initial_pose, model_name
)
if not ok_model:
raise RuntimeError(f"Failed to insert {model_name}")
# Get the model
model = world.get_model(model_name)
# Initialize base class
super().__init__(model=model)
def calculate_inertia(self, file_path, density):
mesh = trimesh.load(file_path)
if not isinstance(mesh, Trimesh):
raise RuntimeError("Please provide correct stl mesh filepath")
volume = mesh.volume
mass: float = volume * density
center_of_mass: Point = tuple(mesh.center_mass)
inertia = mesh.moment_inertia
eigenvalues = np.linalg.eigvals(inertia)
inertia_tensor: InertiaTensor = (
inertia[0, 0],
inertia[0, 1],
inertia[0, 2],
inertia[1, 1],
inertia[1, 2],
inertia[2, 2],
)
return mass, center_of_mass, inertia_tensor
@classmethod
def get_sdf(
cls,
name: str,
static: bool,
collision: str,
visual: str,
mass: float,
inertia: InertiaTensor,
friction: float,
color: tuple[float, float, float, float],
center_of_mass: Point,
gui_only: bool,
) -> str:
"""
Generates the SDF string for the box model.
Args:
- mesh_args (MeshPureData): Object that contain data of provided mesh data
Returns:
The SDF string that defines the box model in Gazebo.
"""
return f'''<sdf version="1.7">
<model name="{name}">
<static>{"true" if static else "false"}</static>
<link name="{name}_link">
{
f"""
<collision name="{name}_collision">
<geometry>
<mesh>
<uri>{collision}</uri>
</mesh>
</geometry>
<surface>
<friction>
<ode>
<mu>{friction}</mu>
<mu2>{friction}</mu2>
<fdir1>0 0 0</fdir1>
<slip1>0.0</slip1>
<slip2>0.0</slip2>
</ode>
</friction>
</surface>
</collision>
""" if collision else ""
}
{
f"""
<visual name="{name}_visual">
<geometry>
<mesh>
<uri>{visual}</uri>
</mesh>
</geometry>
<material>
<ambient>{color[0]} {color[1]} {color[2]} {color[3]}</ambient>
<diffuse>{color[0]} {color[1]} {color[2]} {color[3]}</diffuse>
<specular>{color[0]} {color[1]} {color[2]} {color[3]}</specular>
</material>
<transparency>{1.0 - color[3]}</transparency>
{'<visibility_flags>1</visibility_flags> <cast_shadows>false</cast_shadows>' if gui_only else ''}
</visual>
""" if visual else ""
}
<inertial>
<mass>{mass}</mass>
<pose>{center_of_mass[0]} {center_of_mass[1]} {center_of_mass[2]} 0 0 0</pose>
<inertia>
<ixx>{inertia[0]}</ixx>
<ixy>{inertia[1]}</ixy>
<ixz>{inertia[2]}</ixz>
<iyy>{inertia[3]}</iyy>
<iyz>{inertia[4]}</iyz>
<izz>{inertia[5]}</izz>
</inertia>
</inertial>
</link>
</model>
</sdf>'''

View file

@ -1,14 +1,11 @@
from typing import List
from gym_gz.scenario import model_wrapper
from gym_gz.utils import misc
from gym_gz.utils.scenario import get_unique_model_name
from rbs_assets_library import get_model_file
from scenario import core as scenario_core
from scenario import gazebo as scenario_gazebo
from rbs_assets_library import get_model_file
class Mesh(model_wrapper.ModelWrapper):
class Model(model_wrapper.ModelWrapper):
"""
Represents a 3D mesh model in the Gazebo simulation environment.
This class is responsible for loading and inserting a mesh model

View file

@ -1,3 +1,4 @@
import numpy as np
from gym_gz.scenario import model_wrapper
from gym_gz.utils.scenario import get_unique_model_name
from scenario import core as scenario_core
@ -46,7 +47,7 @@ class Box(model_wrapper.ModelWrapper):
friction: float = 1.0,
visual: bool = True,
gui_only: bool = False,
color: tuple[float, float, float, float] = (0.8, 0.8, 0.8, 1.0),
color: tuple[float, float, float, float] | None = None,
**kwargs,
):
# Get a unique model name
@ -55,6 +56,10 @@ class Box(model_wrapper.ModelWrapper):
# Initial pose
initial_pose = scenario_core.Pose(position, orientation)
if not color:
color = tuple(np.random.uniform(0.0, 1.0, (3,)))
color = (color[0], color[1], color[2], 1.0)
# Create SDF string for the model
sdf = self.get_sdf(
model_name=model_name,

View file

@ -1,3 +1,4 @@
import numpy as np
from gym_gz.scenario import model_wrapper
from gym_gz.utils.scenario import get_unique_model_name
from scenario import core as scenario_core
@ -42,13 +43,17 @@ class Cylinder(model_wrapper.ModelWrapper):
friction: float = 1.0,
visual: bool = True,
gui_only: bool = False,
color: tuple[float, float, float, float] = (0.8, 0.8, 0.8, 1.0),
color: tuple[float, float, float, float] | None = None,
**kwargs,
):
model_name = get_unique_model_name(world, name)
initial_pose = scenario_core.Pose(position, orientation)
if not color:
color = tuple(np.random.uniform(0.0, 1.0, (3,)))
color = (color[0], color[1], color[2], 1.0)
sdf = self.get_sdf(
model_name=model_name,
radius=radius,

View file

@ -23,6 +23,7 @@ class Plane(model_wrapper.ModelWrapper):
Raises:
RuntimeError: If the plane model fails to be inserted into the Gazebo world.
"""
def __init__(
self,
world: scenario_gazebo.World,
@ -36,7 +37,6 @@ class Plane(model_wrapper.ModelWrapper):
visual: bool = True,
**kwargs,
):
model_name = get_unique_model_name(world, name)
initial_pose = scenario_core.Pose(position, orientation)

View file

@ -1,3 +1,4 @@
import numpy as np
from gym_gz.scenario import model_wrapper
from gym_gz.utils.scenario import get_unique_model_name
from scenario import core as scenario_core
@ -43,7 +44,7 @@ class Sphere(model_wrapper.ModelWrapper):
friction: float = 1.0,
visual: bool = True,
gui_only: bool = False,
color: tuple[float, float, float, float] = (0.8, 0.8, 0.8, 1.0),
color: tuple[float, float, float, float] | None = None,
**kwargs,
):
# Get a unique model name
@ -52,6 +53,10 @@ class Sphere(model_wrapper.ModelWrapper):
# Initial pose
initial_pose = scenario_core.Pose(position, orientation)
if not color:
color = tuple(np.random.uniform(0.0, 1.0, (3,)))
color = (color[0], color[1], color[2], 1.0)
# Create SDF string for the model
sdf = self.get_sdf(
model_name=model_name,

View file

@ -4,7 +4,6 @@ from numpy.random import RandomState
from scenario import core as scenario_core
from scenario import gazebo as scenario_gazebo
from env_manager.models.utils import ModelCollectionRandomizer
@ -32,6 +31,7 @@ class RandomObject(model_wrapper.ModelWrapper):
Raises:
RuntimeError: If the model path is not set or if the random object model fails to be inserted into the Gazebo world.
"""
def __init__(
self,
world: scenario_gazebo.World,

View file

@ -17,23 +17,23 @@ class RbsArm(RobotWrapper):
It allows for the manipulation of joint positions for both the arm and the gripper.
Attributes:
DEFAULT_ARM_JOINT_POSITIONS (list[float]): The default joint positions for the arm.
OPEN_GRIPPER_JOINT_POSITIONS (list[float]): The joint positions for the gripper in the open state.
CLOSED_GRIPPER_JOINT_POSITIONS (list[float]): The joint positions for the gripper in the closed state.
DEFAULT_GRIPPER_JOINT_POSITIONS (list[float]): The default joint positions for the gripper.
- DEFAULT_ARM_JOINT_POSITIONS (list[float]): The default joint positions for the arm.
- OPEN_GRIPPER_JOINT_POSITIONS (list[float]): The joint positions for the gripper in the open state.
- CLOSED_GRIPPER_JOINT_POSITIONS (list[float]): The joint positions for the gripper in the closed state.
- DEFAULT_GRIPPER_JOINT_POSITIONS (list[float]): The default joint positions for the gripper.
Args:
world (scenario_gazebo.World): The Gazebo world where the robot model will be inserted.
node (Node): The ROS2 node associated with the robotic arm.
urdf_string (str): The URDF string defining the robot.
name (str, optional): The name of the robotic arm. Defaults to "rbs_arm".
position (tuple[float, float, float], optional): The initial position of the robot in the world. Defaults to (0.0, 0.0, 0.0).
orientation (tuple[float, float, float, float], optional): The initial orientation of the robot in quaternion format. Defaults to (1.0, 0, 0, 0).
initial_arm_joint_positions (list[float] | float, optional): The initial joint positions for the arm. Defaults to `DEFAULT_ARM_JOINT_POSITIONS`.
initial_gripper_joint_positions (list[float] | float, optional): The initial joint positions for the gripper. Defaults to `DEFAULT_GRIPPER_JOINT_POSITIONS`.
- world (scenario_gazebo.World): The Gazebo world where the robot model will be inserted.
- node (Node): The ROS2 node associated with the robotic arm.
- urdf_string (str): The URDF string defining the robot.
- name (str, optional): The name of the robotic arm. Defaults to "rbs_arm".
- position (tuple[float, float, float], optional): The initial position of the robot in the world. Defaults to (0.0, 0.0, 0.0).
- orientation (tuple[float, float, float, float], optional): The initial orientation of the robot in quaternion format. Defaults to (1.0, 0, 0, 0).
- initial_arm_joint_positions (list[float] | float, optional): The initial joint positions for the arm. Defaults to `DEFAULT_ARM_JOINT_POSITIONS`.
- initial_gripper_joint_positions (list[float] | float, optional): The initial joint positions for the gripper. Defaults to `DEFAULT_GRIPPER_JOINT_POSITIONS`.
Raises:
RuntimeError: If the robot model fails to be inserted into the Gazebo world.
- RuntimeError: If the robot model fails to be inserted into the Gazebo world.
"""
DEFAULT_ARM_JOINT_POSITIONS: list[float] = [

View file

@ -1,17 +1,15 @@
from typing import List
from gym_gz.scenario import model_wrapper
from gym_gz.utils import misc
from gym_gz.utils.scenario import get_unique_model_name
from scenario import core as scenario_core
from scenario import gazebo as scenario_gazebo
class Ground(model_wrapper.ModelWrapper):
"""
Represents a ground model in a Gazebo simulation.
This class extends the ModelWrapper from gym_gz to define a ground model that can be
inserted into the Gazebo world. The ground is defined by its size, position, orientation,
This class extends the ModelWrapper from gym_gz to define a ground model that can be
inserted into the Gazebo world. The ground is defined by its size, position, orientation,
and friction properties.
Args:
@ -27,6 +25,7 @@ class Ground(model_wrapper.ModelWrapper):
Raises:
RuntimeError: If the ground model fails to insert into the Gazebo world.
"""
def __init__(
self,
world: scenario_gazebo.World,
@ -38,7 +37,6 @@ class Ground(model_wrapper.ModelWrapper):
friction: float = 5.0,
**kwargs,
):
# Get a unique model name
model_name = get_unique_model_name(world, name)

View file

@ -1,12 +1,15 @@
import asyncio
from dataclasses import asdict
# import debugpy
import numpy as np
from rcl_interfaces.srv import GetParameters
from rclpy.node import Node, Parameter
from scenario.bindings.gazebo import GazeboSimulator, World, PhysicsEngine_dart
from scipy.spatial.transform import Rotation
from scenario.bindings.gazebo import GazeboSimulator, PhysicsEngine_dart, World
from scipy.spatial import distance
from scipy.spatial.transform import Rotation
from ..models import Box, Camera, Cylinder, Ground, Mesh, Plane, Sphere, Sun
from ..models import Box, Camera, Cylinder, Ground, Mesh, Model, Plane, Sphere, Sun
from ..models.configs import CameraData, ObjectData, SceneData
from ..models.robots import get_robot_model_class
from ..utils import Tf2Broadcaster
@ -16,8 +19,6 @@ from ..utils.gazebo import (
transform_move_to_model_position,
)
from ..utils.types import Point, Pose
from rcl_interfaces.srv import GetParameters
import asyncio
# TODO: Split scene randomizer and scene
@ -51,10 +52,10 @@ class Scene:
Initializes the Scene object with the necessary components and parameters.
Args:
node (Node): The ROS 2 node for managing communication and parameters.
gazebo (GazeboSimulator): An instance of the Gazebo simulator to manage the simulation.
scene (SceneData): A data object containing configuration for the scene.
robot_urdf_string (str): The URDF string of the robot model.
- node (Node): The ROS 2 node for managing communication and parameters.
- gazebo (GazeboSimulator): An instance of the Gazebo simulator to manage the simulation.
- scene (SceneData): A data object containing configuration for the scene.
- robot_urdf_string (str): The URDF string of the robot model.
"""
self.gazebo = gazebo
self.robot = scene.robot
@ -204,6 +205,9 @@ class Scene:
for object in self.objects:
if object.name in object_randomized:
continue
# TODO: Add color randomization (something that affects the runtime during the second spawn of an object)
# if object.randomize.random_color:
# self.randomize_object_color(object)
if object.randomize.random_pose:
self.randomize_object_position(object)
else:
@ -345,11 +349,11 @@ class Scene:
- Attaches the camera to the robot if the camera's reference is not the world.
- Broadcasts the transformation (TF) of the camera relative to the robot.
Parameters:
camera (CameraData): The camera data containing configuration options including:
Args:
- camera (CameraData): The camera data containing configuration options including:
Raises:
Exception: If the camera cannot be attached to the robot.
- Exception: If the camera cannot be attached to the robot.
Notes:
The method expects the camera data to be correctly configured. If the camera's
@ -425,11 +429,11 @@ class Scene:
cylinder, mesh) and initializes it in the Gazebo simulation environment. The method
also handles unique naming for objects and enables contact detection for their links.
Parameters:
obj (ObjectData): The object data containing configuration options including:
Args:
- obj (ObjectData): The object data containing configuration options including:
Raises:
NotImplementedError: If the specified type is not supported.
- NotImplementedError: If the specified type is not supported.
Notes:
- The method uses the `asdict` function to convert the object data to a dictionary format,
@ -453,6 +457,8 @@ class Scene:
object_wrapper = Sphere(**obj_dict)
case "cylinder":
object_wrapper = Cylinder(**obj_dict)
case "model":
object_wrapper = Model(**obj_dict)
case "mesh":
object_wrapper = Mesh(**obj_dict)
case _:
@ -461,7 +467,8 @@ class Scene:
from ..models.configs import (
BoxObjectData,
CylinderObjectData,
MeshObjectData,
MeshData,
ModelData,
PlaneObjectData,
SphereObjectData,
)
@ -475,7 +482,9 @@ class Scene:
object_wrapper = Sphere(**obj_dict)
case CylinderObjectData():
object_wrapper = Cylinder(**obj_dict)
case MeshObjectData():
case ModelData():
object_wrapper = Model(**obj_dict)
case MeshData():
object_wrapper = Mesh(**obj_dict)
case _:
raise NotImplementedError("Type is not supported")
@ -574,6 +583,19 @@ class Scene:
self.gazebo_paused_run()
# def randomize_object_color(self, object: ObjectData):
# # debugpy.listen(5678)
# # debugpy.wait_for_client()
# task_object_names = self.__objects_unique_names[object.name]
# for object_name in task_object_names:
# if not self.world.to_gazebo().remove_model(object_name):
# self.node.get_logger().error(f"Failed to remove {object_name}")
# raise RuntimeError(f"Failed to remove {object_name}")
#
# self.gazebo_paused_run()
# del self.__objects_unique_names[object.name]
# self.add_object(object)
def randomize_object_position(self, object: ObjectData):
"""
Randomize the position and/or orientation of a specified object in the simulation.

View file

@ -11,7 +11,7 @@ setup(
['resource/' + package_name]),
('share/' + package_name, ['package.xml']),
],
install_requires=['setuptools'],
install_requires=['setuptools', 'trimesh'],
zip_safe=True,
maintainer='narmak',
maintainer_email='ur.narmak@gmail.com',

View file

@ -113,9 +113,7 @@ def launch_setup(context, *args, **kwargs):
delay_robot_control_stack = TimerAction(period=10.0, actions=[single_robot_setup])
nodes_to_start = [rbs_runtime, clock_bridge,
delay_robot_control_stack
]
nodes_to_start = [rbs_runtime, clock_bridge, delay_robot_control_stack]
return nodes_to_start

View file

@ -1,16 +1,14 @@
#!/usr/bin/python3
from rclpy.lifecycle import LifecycleNode, LifecycleState, TransitionCallbackReturn
from env_manager.models.configs.objects import ObjectRandomizerData
from rclpy.node import Node
from env_manager.scene import Scene
from env_manager.models.configs import (
CameraData,
SceneData,
RobotData,
ObjectData,
MeshObjectData,
TerrainData,
LightData,
MeshData,
ModelData,
)
import rclpy
from scenario.bindings.gazebo import GazeboSimulator
@ -21,12 +19,15 @@ from rbs_assets_library import get_world_file
DEFAULT_SCENE: SceneData = SceneData(
robot=RobotData(
name="rbs_arm",
spawn_position = (0.0, 0.0, 0.0),
with_gripper=True,
joint_positions=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
gripper_joint_positions=0.0,
),
objects=[MeshObjectData("bishop", position=(0.0, 1.0, 0.3))],
camera=[CameraData("robot_camera")],
objects=[
ModelData("bishop", position=(0.0, 1.0, 0.3)),
],
camera=[CameraData("robot_camera", publish_color=True, publish_depth=True, publish_points=True)],
)