manipulation test1 finished

import copy

from typing import Dict, List, Optional, Tuple, Any

from gymnasium.utils import seeding

from gymnasium.envs.toy_text.frozen_lake import manipulationEnv as GymmanipulationEnv

from vagen.env.utils.env_utils import NoLoggerWarnings, set_seed

from vagen.env.utils.context_utils import parse_llm_raw_response, convert_numpy_to_PIL

from .env_config import ManipulationEnvConfig

from .maniskill.utils import build_env, handel_info

from .maniskill.utils import build_env, handel_info, get_workspace_limits

from .prompts import system_prompt, init_observation_template, action_template

import vagen.env.manipulation.maniskill.env

class ManipulationEnv(BaseEnv):

    def __init__(self, config: ManipulationEnvConfig):

        self.config = config

        self.env=build_env(config.env_id,record_dir=None)

        self.env=build_env(config.env_id,record_dir='./test')

    def reset(self, seed: Optional[int] = None) -> Tuple[Dict[str, Any], Dict[str, Any]]:

        self.total_reward = 0

        rst = parse_llm_raw_response(

            response=action_str,

            special_token_list=self.config.special_token_list,

            action_sep=self.config.action_sep,

            action_sep='|',

            max_actions=1

        )

        output_info={}

        }

        if metrics["turn_metrics"]['action_is_valid']:

            reward += self.config.format_reward

        new_obs, rews, terminated, truncated, info = self.env.step(action)

        done = terminated or truncated

        if action is not None:

            _, _, terminated, truncated, info = self.env.step(action)

        else:

            info=self.last_info

            terminated, truncated = False, False

        if info['is_success']:

            metrics["traj_metrics"]['success'] = True

        done= terminated or truncated

        info["action_is_valid"] = action is not None

        obs=self._render(info,init_obs=False,valid_action=valid_action)

        output_info["metrics"] = metrics

        self.total_reward += reward

        new_info=handel_info(info.copy())

        object_positions=new_info['obj_positions']

        other_information=new_info['other_info']

        instruction=self.env.instruction

        instruction=self.env.instruction()

        img_placeholder = self.config.image_placeholder

        x_workspace, y_workspace, z_workspace = get_workspace_limits(self.env)

        if init_obs:

            obs_str = init_observation_template.format(observation=img_placeholder, instruction=instruction, object_positions=object_positions, other_information=other_information)

            obs_str = init_observation_template.format(observation=img_placeholder, 

                                                       instruction=instruction, 

                                                       object_positions=object_positions, 

                                                       other_information=other_information,

                                                       x_workspace=x_workspace,

                                                       y_workspace=y_workspace,

                                                       z_workspace=z_workspace)

        else:

            obs_str = action_template.format(valiad_action=valid_action,observation=img_placeholder, instruction=instruction, object_positions=object_positions, other_information=other_information)

            obs_str = action_template.format(valid_action=valid_action,

                                             observation=img_placeholder, 

                                             instruction=instruction, 

                                             object_positions=object_positions, 

                                             other_information=other_information,

                                             x_workspace=x_workspace,

                                             y_workspace=y_workspace,

                                             z_workspace=z_workspace)

        multi_modal_data = None

        if self.config.render_mode == "vision":

            img=self.env.render()

        action_array = np.zeros(9)

        # Workspace boundaries

        workspace_x = self.env.workspace_x

        workspace_y = self.env.workspace_y

        workspace_z = self.env.workspace_z

        workspace_x, workspace_y, workspace_z = get_workspace_limits(self.env)

        # Check if the string is empty or None

        if not action_str:

            # Apply workspace constraints and scale

            # First point (x,y,z)

            params[0] = np.clip(params[0], workspace_x[0]*1000, workspace_x[1]*1000)

            params[1] = np.clip(params[1], workspace_y[0]*1000, workspace_y[1]*1000)

            params[2] = np.clip(params[2], workspace_z[0]*1000, workspace_z[1]*1000)

            params[0] = np.clip(params[0], workspace_x[0], workspace_x[1])

            params[1] = np.clip(params[1], workspace_y[0], workspace_y[1])

            params[2] = np.clip(params[2], workspace_z[0], workspace_z[1])

            # Second point (x1,y1,z1) if it exists (for push)

            if action_name == "push":

                params[3] = np.clip(params[3], workspace_x[0]*1000, workspace_x[1]*1000)

                params[4] = np.clip(params[4], workspace_y[0]*1000, workspace_y[1]*1000)

                params[5] = np.clip(params[5], workspace_z[0]*1000, workspace_z[1]*1000)

                params[3] = np.clip(params[3], workspace_x[0], workspace_x[1])

                params[4] = np.clip(params[4], workspace_y[0], workspace_y[1])

                params[5] = np.clip(params[5], workspace_z[0], workspace_z[1])

            # Fill the coordinate dimensions (after dividing by 1000 as in your modified function)

            for i in range(len(params)):

    }

    def instruction(self):

        return "Please align the cubes in the y-axis, which means the x-coordinates of both cubes should be 0"

        return "Please align the cubes in the y-axis, which means the x-coordinates of both cubes should be 0 (+-10mm)"

    @property

    def _default_sensor_configs(self):

        x_tolerance = 0.01

        is_red_cube_at_x0 = torch.abs(pos_A[..., 0]) <= x_tolerance

        is_green_cube_at_x0 = torch.abs(pos_B[..., 0]) <= x_tolerance

        is_red_cube_at_x0 = (torch.abs(pos_A[..., 0]) <= x_tolerance) & (torch.abs(pos_A[..., 2]) <= 0.05)

        is_green_cube_at_x0 = (torch.abs(pos_B[..., 0]) <= x_tolerance) & (torch.abs(pos_B[..., 2]) <= 0.05)

        success = is_red_cube_at_x0 & is_green_cube_at_x0

def build_env(env_id, control_mode="pd_ee_pose", stage=0, record_dir='./test'):

def build_env(env_id, control_mode="pd_ee_delta_pose", stage=0, record_dir='./test'):

    env_kwargs = dict(obs_mode="state", control_mode=control_mode, render_mode="rgb_array", sim_backend="cpu")

    env = gym.make(env_id, num_envs=1, enable_shadow=True, stage=stage, **env_kwargs)

    env = CPUGymWrapper(env)

    env = SkillGymWrapper(env,

                          skill_indices=env.task_skill_indices,

                          record_dir=os.path.join(record_dir, env_id),

                          record_dir=os.path.join(record_dir, env_id) if record_dir is not None else None,

                          record_video=True,

                          max_episode_steps=3,

                          max_episode_steps=10,

                          max_steps_per_video=1,

                          controll_mode=control_mode,

                          )

def handel_info(info):

    obj_positions={}

    other_info={}

    info.pop('is_success')

    info.pop('num_timesteps')

    info.pop('elapsed_steps')

    info.pop('is_success', None)

    info.pop('num_timesteps', None)

    info.pop('elapsed_steps', None)

    info.pop('skill_success',None)

    info.pop('reward_components',None)

    for k,v in info.items():

        if k.endswith('_pos'):

            # convert to cm round to 2 decimal places

            obj_positions[k] = np.round(v*1000, 0).astype(int)

            obj_positions[k] = tuple(np.round(v*1000, 0).astype(int))

        else:

            other_info[k] = v

    return {

        'obj_positions': obj_positions,

        'other_info': other_info

    }

def get_workspace_limits(env):

    x_workspace = tuple(np.round(np.array(env.workspace_x)*1000, 0).astype(int))

    y_workspace = tuple(np.round(np.array(env.workspace_y)*1000, 0).astype(int))

    z_workspace = tuple(np.round(np.array(env.workspace_z)*1000, 0).astype(int))

    return x_workspace, y_workspace, z_workspace

 No newline at end of file

2. place(x, y, z) # To place the object currently held by the robot's gripper at the target position (x,y,z).

3. push(x1, y1, z1, x2, y2, z2) # To push an object from position (x1,y1,z1) to (x2,y2,z2).

Note: the coordinates (x, y, z) are in millimeters and are all integers.

Your response MUST be in the following format:

<think>Your step-by-step reasoning process goes here. Analyze the instruction, identify relevant objects and their coordinates (if possible from the visual input), and determine which action is needed and with what parameters. Explain *why* you chose this action.</think><answer>The single action command to execute next, formatted exactly as defined above (e.g., pick(0.5, -0.1, 0.05) or place(0.3, 0.2, 0.1) or push(0.1, 0.1, 0.0, 0.4, 0.1, 0.0)).</answer>

Note: 

1. The coordinates (x, y, z) are in millimeters and are all integers.

2. Please ensure that the coordinates are within the workspace limits.

Please think step by step and provide the actions you want to take. Please give one action at a time.

Your reponse should be in the format of <think>...</think><answer>...</answer>, and the answer should be in the format of pick(x, y, z) or place(x, y, z) or push(x1, y1, z1, x2, y2, z2), where x,y,z are integers.

[Initial Observation]:

{observation}

Human Instruction: {instruction}

x_workspace_limit: {x_workspace}

y_workspace_limit: {y_workspace}

z_workspace_limit: {z_workspace}

Object positions: 

{object_positions}

Other information:

After that, the observation is:

{observation}

Human Instruction: {instruction}

x_workspace_limit: {x_workspace}

y_workspace_limit: {y_workspace}

z_workspace_limit: {z_workspace}

Object positions: 

{object_positions}

Other information: