Implementation:ARISE Initiative Robosuite NutAssembly Environment
| Knowledge Sources | |
|---|---|
| Domains | Robotics, Simulation, Manipulation |
| Last Updated | 2026-02-15 07:00 GMT |
Overview
Concrete tool for simulating a nut-on-peg assembly manipulation task provided by robosuite.
Description
The NutAssembly class implements a single-arm robot nut assembly task in which one or more nuts must be placed onto their corresponding pegs on a tabletop workspace. The environment features two nut types -- a square nut and a round nut -- each with a dedicated peg. The robot must pick up each nut, lift it, and accurately place it around the correct peg. This task tests grasping, lifting, and precise placement skills.
The reward function operates in both sparse and dense modes. In sparse mode, a discrete reward of 1.0 is provided per nut correctly placed on its peg. In dense (shaped) mode, four staged reward components are computed: reaching (in [0, 0.1], based on distance to closest nut), grasping (in {0, 0.35}, binary for contact with any active nut), lifting (in {0, [0.35, 0.5]}, proportional to lifting height), and hovering (in {0, [0.5, 0.7]}, proportional to distance from nut to its target peg). Successfully placed nuts always contribute 1.0 to the total reward. The final reward is normalized and scaled by reward_scale / 2.0 (or reward_scale / 1.0 if only a single nut is used).
The environment supports three task modes via the single_object_mode parameter: mode 0 uses both nuts (full task), mode 1 randomly selects one nut type per reset, and mode 2 fixes the nut type to a specified value. Convenience subclasses NutAssemblySingle, NutAssemblySquare, and NutAssemblyRound are provided for the simplified variants. Success is determined when all active nuts are correctly placed on their corresponding pegs while the gripper is not in close proximity.
Usage
Use this environment for training and evaluating single-arm pick-and-place policies that require both object discrimination and precise spatial placement. The multi-object and single-object modes allow progressive curriculum training.
Code Reference
Source Location
- Repository: ARISE_Initiative_Robosuite
- File: robosuite/environments/manipulation/nut_assembly.py
- Lines: 1-711
Signature
class NutAssembly(ManipulationEnv):
def __init__(
self,
robots,
env_configuration="default",
controller_configs=None,
gripper_types="default",
base_types="default",
initialization_noise="default",
table_full_size=(0.8, 0.8, 0.05),
table_friction=(1, 0.005, 0.0001),
table_offset=(0, 0, 0.82),
use_camera_obs=True,
use_object_obs=True,
reward_scale=1.0,
reward_shaping=False,
placement_initializer=None,
single_object_mode=0,
nut_type=None,
has_renderer=False,
has_offscreen_renderer=True,
render_camera="frontview",
render_collision_mesh=False,
render_visual_mesh=True,
render_gpu_device_id=-1,
control_freq=20,
lite_physics=True,
horizon=1000,
ignore_done=False,
hard_reset=True,
camera_names="agentview",
camera_heights=256,
camera_widths=256,
camera_depths=False,
camera_segmentations=None,
renderer="mjviewer",
renderer_config=None,
seed=None,
):
Import
from robosuite.environments.manipulation.nut_assembly import NutAssembly
# Convenience variants:
from robosuite.environments.manipulation.nut_assembly import NutAssemblySingle
from robosuite.environments.manipulation.nut_assembly import NutAssemblySquare
from robosuite.environments.manipulation.nut_assembly import NutAssemblyRound
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| robots | str or list of str | Yes | Specification for a single single-arm robot (e.g., "Panda") |
| single_object_mode | int | No | 0 = both nuts, 1 = random single nut per reset, 2 = fixed single nut. Default: 0 |
| nut_type | str | No | "round" or "square"; only used when single_object_mode=2. Default: None |
| table_full_size | 3-tuple | No | (x, y, z) dimensions of the table. Default: (0.8, 0.8, 0.05) |
| table_friction | 3-tuple | No | MuJoCo friction parameters for the table. Default: (1, 0.005, 0.0001) |
| reward_scale | None or float | No | Scales the normalized reward. Default: 1.0 |
| reward_shaping | bool | No | If True, uses dense staged rewards. Default: False |
| placement_initializer | ObjectPositionSampler | No | Custom placement sampler. Default: SequentialCompositeSampler |
| horizon | int | No | Episode length in timesteps. Default: 1000 |
Outputs
| Name | Type | Description |
|---|---|---|
| {NutName}_pos | np.array (3,) | 3D position of each nut body |
| {NutName}_quat | np.array (4,) | Quaternion orientation (xyzw) of each nut |
| {NutName}_to_{prefix}eef_pos | np.array (3,) | Relative position from nut to end-effector |
| {NutName}_to_{prefix}eef_quat | np.array (4,) | Relative orientation from nut to end-effector |
| nut_id | int | Active nut ID (only in single_object_mode=1) |
| reward | float | Scalar reward value per step |
Usage Examples
import robosuite as suite
import numpy as np
# Create a NutAssembly environment with both nut types
env = suite.make(
env_name="NutAssembly",
robots="Panda",
has_renderer=False,
has_offscreen_renderer=False,
use_camera_obs=False,
use_object_obs=True,
reward_shaping=True,
horizon=1000,
)
obs = env.reset()
for i in range(1000):
action = np.random.randn(env.action_dim)
obs, reward, done, info = env.step(action)
if done:
break
env.close()
# Alternatively, use the simplified square nut variant
env_square = suite.make(
env_name="NutAssemblySquare",
robots="Sawyer",
has_renderer=False,
has_offscreen_renderer=False,
use_camera_obs=False,
)