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

fix spacemouse action transformation direction #38

Merged
merged 4 commits into from
Apr 25, 2024
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
6 changes: 4 additions & 2 deletions examples/async_bin_relocation_fwbw_drq/record_bc_demos.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,8 +75,10 @@ def on_esc(key):
transitions = []
while demos_count < demos_needed:

next_obs, rew, done, truncated, info = env.step(action=np.zeros((7,)))
actions = info["intervene_action"]
actions = np.zeros((7,))
next_obs, rew, done, truncated, info = env.step(action=actions)
if "intervene_action" in info:
actions = info["intervene_action"]

transition = copy.deepcopy(
dict(
Expand Down
6 changes: 4 additions & 2 deletions examples/async_bin_relocation_fwbw_drq/record_demo.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,8 +85,10 @@
len(fw_transitions) < transitions_needed
or len(bw_transitions) < transitions_needed
):
next_obs, rew, done, truncated, info = env.step(action=np.zeros((7,)))
actions = info["intervene_action"]
actions = np.zeros((7,))
next_obs, rew, done, truncated, info = env.step(action=actions)
if "intervene_action" in info:
actions = info["intervene_action"]

transition = copy.deepcopy(
dict(
Expand Down
6 changes: 4 additions & 2 deletions examples/async_bin_relocation_fwbw_drq/record_transitions.py
Original file line number Diff line number Diff line change
Expand Up @@ -98,9 +98,11 @@ def check_all_done():

# Loop until we have enough transitions
while not check_all_done():
next_obs, rew, done, truncated, info = env.step(action=np.zeros((7,)))
actions = np.zeros((7,))
next_obs, rew, done, truncated, info = env.step(action=actions)
if "intervene_action" in info:
actions = info["intervene_action"]
next_obs = env.get_front_cam_obs()
actions = info["intervene_action"]

transition = copy.deepcopy(
dict(
Expand Down
6 changes: 4 additions & 2 deletions examples/async_bin_relocation_fwbw_drq/test_classifier.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,8 +55,10 @@
obs, _ = env.reset()

for i in tqdm(range(1000)):
next_obs, rew, done, truncated, info = env.step(action=np.zeros((7,)))
actions = info["intervene_action"]
actions = np.zeros((7,))
next_obs, rew, done, truncated, info = env.step(action=actions)
if "intervene_action" in info:
actions = info["intervene_action"]

obs = next_obs

Expand Down
6 changes: 4 additions & 2 deletions examples/async_cable_route_drq/record_demo.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,8 +64,10 @@
raise PermissionError(f"No permission to write to {file_dir}")

while success_count < success_needed:
next_obs, rew, done, truncated, info = env.step(action=np.zeros((6,)))
actions = info["intervene_action"]
actions = np.zeros((6,))
next_obs, rew, done, truncated, info = env.step(action=actions)
if "intervene_action" in info:
actions = info["intervene_action"]

transition = copy.deepcopy(
dict(
Expand Down
6 changes: 4 additions & 2 deletions examples/async_cable_route_drq/test_classifier.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,8 +43,10 @@
obs, _ = env.reset()

for i in tqdm(range(1000)):
next_obs, rew, done, truncated, info = env.step(action=np.zeros((6,)))
actions = info["intervene_action"]
actions = np.zeros((6,))
next_obs, rew, done, truncated, info = env.step(action=actions)
if "intervene_action" in info:
actions = info["intervene_action"]

obs = next_obs

Expand Down
6 changes: 4 additions & 2 deletions examples/async_pcb_insert_drq/record_demo.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,8 +49,10 @@
raise PermissionError(f"No permission to write to {file_dir}")

while success_count < success_needed:
next_obs, rew, done, truncated, info = env.step(action=np.zeros((6,)))
actions = info["intervene_action"]
actions = np.zeros((6,))
next_obs, rew, done, truncated, info = env.step(action=actions)
if "intervene_action" in info:
actions = info["intervene_action"]

transition = copy.deepcopy(
dict(
Expand Down
6 changes: 4 additions & 2 deletions examples/async_peg_insert_drq/record_demo.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,8 +49,10 @@
raise PermissionError(f"No permission to write to {file_dir}")

while success_count < success_needed:
next_obs, rew, done, truncated, info = env.step(action=np.zeros((6,)))
actions = info["intervene_action"]
actions = np.zeros((6,))
next_obs, rew, done, truncated, info = env.step(action=actions)
if "intervene_action" in info:
actions = info["intervene_action"]
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

how does this code run correctly without issues previously?

if "intervene_action" is None, the previous code always assumes
actions = info["intervene_action"]

and this doesn't throw out errors?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My bad, there was a change in the IntervenetionWrapper that I missed. Now it puts intervene_action in info only if intervene_action is nonzero

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i was more talking the code in the left

        next_obs, rew, done, truncated, info = env.step(action=np.zeros((7,)))
        actions = info["intervene_action"]

does this code run without issues previously? the actions can be unassigned in certain cases

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it used to be always assigned. If no intervention was provided, then the policy action is just put in info['intervene_action'], but this meant the intervene_action were in two different frames depending on whether it's from the policy or the spacemouse


transition = copy.deepcopy(
dict(
Expand Down
17 changes: 15 additions & 2 deletions serl_robot_infra/franka_env/envs/relative_env.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,9 @@ class RelativeFrame(gym.Wrapper):
}
),
......
}, and at least 6 DoF action space with (x, y, z, rx, ry, rz, ...)
}, and at least 6 DoF action space with (x, y, z, rx, ry, rz, ...).
By convention, the 7th dimension of the action space is used for the gripper.

"""

def __init__(self, env: Env, include_relative_pose=True):
Expand All @@ -44,7 +46,9 @@ def step(self, action: np.ndarray):

# this is to convert the spacemouse intervention action
if "intervene_action" in info:
info["intervene_action"] = self.transform_action(info["intervene_action"])
info["intervene_action"] = self.transform_action_inv(
info["intervene_action"]
)

# Update adjoint matrix
self.adjoint_matrix = construct_adjoint_matrix(obs["state"]["tcp_pose"])
Expand Down Expand Up @@ -94,3 +98,12 @@ def transform_action(self, action: np.ndarray):
action = np.array(action) # in case action is a jax read-only array
action[:6] = self.adjoint_matrix @ action[:6]
return action

def transform_action_inv(self, action: np.ndarray):
"""
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you add description that this assumes last dimension is gripper open/close?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There's actually description in the docstring for the entire wrapper class. It just requires the first 6 dimensions to be translation+rotation, and does not assume anything about gripper dimension.

class RelativeFrame(gym.Wrapper):
    """
    This wrapper transforms the observation and action to be expressed in the end-effector frame.
    Optionally, it can transform the tcp_pose into a relative frame defined as the reset pose.

    This wrapper is expected to be used on top of the base Franka environment, which has the following
    observation space:
    {
        "state": spaces.Dict(
            {
                "tcp_pose": spaces.Box(-np.inf, np.inf, shape=(7,)), # xyz + quat
                ......
            }
        ),
        ......
    }, and at least 6 DoF action space with (x, y, z, rx, ry, rz, ...)
    """

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok, then what's the assumption for the 7th>?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Doesn't assume anything for the 7th. It can also be non-existant. All the code only indices the first 6 dimensions action[:6]

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

then can we add a simple description of the expected 7th dimension? e.g., discrete/continuous gripper status

Transform action from spatial(base) frame into body(end-effector) frame
using the adjoint matrix.
"""
action = np.array(action)
action[:6] = np.linalg.inv(self.adjoint_matrix) @ action[:6]
return action
11 changes: 7 additions & 4 deletions serl_robot_infra/franka_env/envs/wrappers.py
Original file line number Diff line number Diff line change
Expand Up @@ -202,14 +202,17 @@ def action(self, action: np.ndarray) -> np.ndarray:
expert_a = np.concatenate((expert_a, gripper_action), axis=0)

if time.time() - self.last_intervene < 0.5:
return expert_a
return expert_a, True

return action
return action, False

def step(self, action):
new_action = self.action(action)

new_action, replaced = self.action(action)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

replaced is always False?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's set to true in line 205, which is when the current time is within 0.5 seconds of the previous nonzero spacemouse action (0.5 sec delay until the control is handed back to the policy)


obs, rew, done, truncated, info = self.env.step(new_action)
info["intervene_action"] = new_action
if replaced:
info["intervene_action"] = new_action
info["left"] = self.left
info["right"] = self.right
return obs, rew, done, truncated, info
Loading