Praxis / cliport /generated_tasks /build_spiral_fixed.py
leofeltrin's picture
Melhorar espiral: aumentar raio (0.05-0.20) e usar 1.5 rotacoes (3*pi) para melhor visibilidade
7d676ca
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
class BuildSpiralWithColoredBlocks(Task):
"""Build a spiral with colored blocks starting from the center."""
def __init__(self):
super().__init__()
self.max_steps = 10
self.lang_template = "build a spiral with the {blocks}"
self.task_completed_desc = "done building spiral."
self.additional_reset()
def reset(self, env):
super().reset(env)
# Define block colors.
colors = [
utils.COLORS['red'], utils.COLORS['orange'], utils.COLORS['yellow'],
utils.COLORS['green'], utils.COLORS['blue'], utils.COLORS['purple']
]
# Add blocks.
block_size = (0.04, 0.04, 0.04)
block_urdf = 'stacking/block.urdf'
objs = []
for i in range(6):
block_pose = self.get_random_pose(env, block_size)
block_id = env.add_object(block_urdf, block_pose, color=colors[i])
objs.append(block_id)
# Define spiral placement using polar coordinates with increasing radius
# Spiral: radius increases, angle increases (1.5 full rotations = 3*pi for better visibility)
num_blocks = 6
spiral_center = (0.5, 0.0, 0.02) # Center of workspace, fixed pose
base_rotation = (0, 0, 0, 1)
base_pose = (spiral_center, base_rotation)
# Spiral parameters:
# - Radius increases from 0.05 to 0.20 (larger range for better visibility)
# - Angle goes from 0 to 3*pi (1.5 full rotations for clearer spiral pattern)
spiral_place_pos = []
for i in range(num_blocks):
# Angle increases from 0 to 3*pi (1.5 full rotations)
angle = (i / (num_blocks - 1)) * 3 * np.pi if num_blocks > 1 else 0
# Radius increases linearly from 0.05 to 0.20 (larger spacing)
radius = 0.05 + (i / (num_blocks - 1)) * 0.15 if num_blocks > 1 else 0.05
# Convert to Cartesian coordinates relative to center
x = radius * np.cos(angle)
y = radius * np.sin(angle)
z = 0.02 # All blocks at same height for 2D spiral
spiral_place_pos.append((x, y, z))
# Apply base pose to all spiral positions
targs = [(utils.apply(base_pose, pos), base_pose[1]) for pos in spiral_place_pos]
# Goal: blocks are placed in a spiral pattern.
language_goal = self.lang_template.format(blocks="red, orange, yellow, green, blue and purple blocks")
self.add_goal(objs=objs, matches=np.eye(6), targ_poses=targs, replace=False,
rotations=True, metric='pose', params=None, step_max_reward=1 / 6,
symmetries=[np.pi/2]*6, language_goal=language_goal)