Skip to content
This repository was archived by the owner on Oct 4, 2023. It is now read-only.

Commit 01a9e48

Browse files
twahrburg-googletwahrburg-google
committed
Pyreach sync 20220314
1 parent 00ddd4e commit 01a9e48

File tree

13 files changed

+326
-37
lines changed

13 files changed

+326
-37
lines changed

‎pyreach/arm.py‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,9 @@ class ArmState:
102102
device_type: The device type of the arm..
103103
device_name: The device name of the arm.
104104
joint_angles: The arm joint angles in radians.
105-
pose: The pose of the arm.
105+
pose: The pose of the arm (alias for tcp_t_base).
106+
tcp_t_base: The pose of the tool center point (TCP) relative to the base of
107+
the arm.
106108
force: The arm joint force values in Newton-meters. Robot specific.
107109
is_protective_stopped: True if the robot is protective stopped.
108110
is_emergency_stopped: True if the robot is emergency stopped.
@@ -112,6 +114,9 @@ class ArmState:
112114
is_program_running: True if a program is running on the robot.
113115
is_robot_power_on: True if the robot power is on.
114116
robot_mode: Mode of operation of the robot.
117+
tip_adjust_t_tcp: Adjust transform for the tooltip.
118+
tip_adjust_t_base: Pose with the tip adjust applied relative to the base of
119+
the arm.
115120
"""
116121

117122
time: float = 0.0
@@ -122,6 +127,9 @@ class ArmState:
122127
pose: core.Pose = core.Pose(
123128
core.Translation(0.0, 0.0, 0.0), core.AxisAngle.from_tuple(
124129
(0.0, 0.0, 0.0)))
130+
tcp_t_base: core.Pose = core.Pose(
131+
core.Translation(0.0, 0.0, 0.0), core.AxisAngle.from_tuple(
132+
(0.0, 0.0, 0.0)))
125133
force: Tuple[float, ...] = (0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
126134
is_protective_stopped: bool = False
127135
is_emergency_stopped: bool = False
@@ -131,6 +139,8 @@ class ArmState:
131139
is_program_running: bool = False
132140
is_robot_power_on: bool = False
133141
robot_mode: RobotMode = RobotMode.DEFAULT
142+
tip_adjust_t_tcp: Optional[core.Pose] = None
143+
tip_adjust_t_base: Optional[core.Pose] = None
134144

135145

136146
class Arm(object):

‎pyreach/factory.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -179,7 +179,7 @@ class LocalPlaybackHostFactory(HostFactory):
179179
_working_directory: str
180180
_robot_id: str
181181
_client_id: Optional[str]
182-
_select_client: bool
182+
_select_client_id: bool
183183
_snapshot_run_id: Optional[str]
184184
_select_snapshot_run_id: bool
185185
_kwargs: Any
@@ -188,7 +188,7 @@ def __init__(self,
188188
robot_id: str,
189189
working_directory: str,
190190
client_id: Optional[str] = None,
191-
select_client: bool = False,
191+
select_client_id: bool = False,
192192
snapshot_run_id: Optional[str] = None,
193193
select_snapshot_run_id: bool = False,
194194
**kwargs: Any) -> None:
@@ -198,7 +198,7 @@ def __init__(self,
198198
robot_id: the robot id to connect to.
199199
working_directory: optional directory to run within.
200200
client_id: the client ID to simulate.
201-
select_client: if client ID is None, will select first client.
201+
select_client_id: if client ID is None, will select first client.
202202
snapshot_run_id: if specified, will select the given snapshot by
203203
gym_run_id.
204204
select_snapshot_run_id: if specified, but snapshot_run_id is None, will
@@ -208,7 +208,7 @@ def __init__(self,
208208
self._robot_id = robot_id
209209
self._working_directory = working_directory
210210
self._client_id = client_id
211-
self._select_client = select_client
211+
self._select_client_id = select_client_id
212212
self._snapshot_run_id = snapshot_run_id
213213
self._select_snapshot_run_id = select_snapshot_run_id
214214
if kwargs is None:
@@ -228,7 +228,7 @@ def connect(self) -> pyreach.Host:
228228

229229
fn = getattr(mod, 'connect_logs_directory')
230230
return fn(self._robot_id, self._working_directory, self._client_id,
231-
self._select_client, self._snapshot_run_id,
231+
self._select_client_id, self._snapshot_run_id,
232232
self._select_snapshot_run_id, self._kwargs)
233233

234234

‎pyreach/gyms/reach_env.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -228,6 +228,8 @@ def __init__(self,
228228
if not host:
229229
if not connection_string:
230230
connection_string = ""
231+
else:
232+
connection_string = "select-snapshot-run-id=true," + connection_string
231233
host = factory.ConnectionFactory(
232234
connection_string=connection_string, **host_kwargs).connect()
233235

‎pyreach/impl/arm_impl.py‎

Lines changed: 23 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,7 @@ def fk(self, joints: List[float]) -> Optional[List[float]]:
9090
the pose.
9191
"""
9292
pose = self._ik.fk(joints)
93-
if pose:
93+
if pose is not None:
9494
return pose.tolist()
9595
return None
9696

@@ -1294,7 +1294,8 @@ def get_message_supplement(
12941294
self._supported_controllers = tuple(descs)
12951295
if (self._device_name == msg.device_name and msg.device_type == "robot" and
12961296
msg.data_type == "robot-state"):
1297-
return self._arm_state_from_message(self._arm_type, msg)
1297+
return self._arm_state_from_message(self._arm_type, msg,
1298+
self._update_data_cache())
12981299
return None
12991300

13001301
def get_key_values(self) -> Set[device_base.KeyValueKey]:
@@ -1532,13 +1533,15 @@ def get_wrapper(
15321533
return (self, tuple(self._internal_devices), ArmImpl(self))
15331534

15341535
@classmethod
1535-
def _arm_state_from_message(cls, arm_type: arm.ArmType,
1536-
msg: types_gen.DeviceData) -> "arm.ArmState":
1536+
def _arm_state_from_message(
1537+
cls, arm_type: arm.ArmType, msg: types_gen.DeviceData,
1538+
data_cache: Optional[_ArmDataCache]) -> "arm.ArmState":
15371539
"""Return the State from a device data message.
15381540
15391541
Args:
15401542
arm_type: The type of the arm.
15411543
msg: The device data to decode.
1544+
data_cache: The arm data cache.
15421545
15431546
Returns:
15441547
The Arm state.
@@ -1576,13 +1579,24 @@ def _arm_state_from_message(cls, arm_type: arm.ArmType,
15761579
robot_mode = arm.RobotMode.from_string(msg.robot_mode)
15771580
except ValueError:
15781581
robot_mode = arm.RobotMode.DEFAULT
1582+
1583+
adjust_pose: Optional[core.Pose] = None
1584+
tip_adjust_transform: Optional[core.Pose] = None
1585+
if data_cache and data_cache.tip_adjust_transform is not None:
1586+
adjusted_pose = transform_util.multiply_pose(
1587+
np.array(pose, dtype=np.float64),
1588+
data_cache.tip_adjust_transform).tolist()
1589+
adjust_pose = core.Pose.from_list(adjusted_pose)
1590+
tip_adjust_transform = core.Pose.from_list(
1591+
data_cache.tip_adjust_transform.tolist())
1592+
pose_t = core.Pose.from_list(pose)
15791593
return arm.ArmState(
15801594
utils.time_at_timestamp(msg.ts), msg.seq,
1581-
msg.device_type, msg.device_name, tuple(joints),
1582-
core.Pose.from_list(pose), tuple(force), msg.is_protective_stopped,
1583-
msg.is_emergency_stopped, msg.is_safeguard_stopped, msg.is_reduced_mode,
1584-
msg.safety_message, msg.is_program_running, msg.is_robot_power_on,
1585-
robot_mode)
1595+
msg.device_type, msg.device_name, tuple(joints), pose_t, pose_t,
1596+
tuple(force), msg.is_protective_stopped, msg.is_emergency_stopped,
1597+
msg.is_safeguard_stopped, msg.is_reduced_mode, msg.safety_message,
1598+
msg.is_program_running, msg.is_robot_power_on, robot_mode,
1599+
tip_adjust_transform, adjust_pose)
15861600

15871601
@property
15881602
def arm_type(self) -> arm.ArmType:

‎pyreach/impl/arm_impl_test.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -757,11 +757,33 @@ def _verify_state(self, state: arm.ArmState) -> None:
757757
self.assertEqual(state.pose.orientation.axis_angle.rx, 0.1186541477907012)
758758
self.assertEqual(state.pose.orientation.axis_angle.ry, -3.075739605978813)
759759
self.assertEqual(state.pose.orientation.axis_angle.rz, -0.01666054968029903)
760+
self.assertEqual(state.tcp_t_base.position.x, 0.1961122234076476)
761+
self.assertEqual(state.tcp_t_base.position.y, -0.7146550885497088)
762+
self.assertEqual(state.tcp_t_base.position.z, 0.1642837633972083)
763+
self.assertEqual(state.tcp_t_base.orientation.axis_angle.rx,
764+
0.1186541477907012)
765+
self.assertEqual(state.tcp_t_base.orientation.axis_angle.ry,
766+
-3.075739605978813)
767+
self.assertEqual(state.tcp_t_base.orientation.axis_angle.rz,
768+
-0.01666054968029903)
760769
self.assertEqual(
761770
state.force,
762771
(-6.514813773909755, 5.273554158068752, -3.661764071560078,
763772
-0.1562602891054569, -0.2400126816079334, 0.1363557378946311))
764773
self.assertFalse(state.is_protective_stopped)
774+
tip_adjust = state.tip_adjust_t_tcp
775+
self.assertIsNotNone(tip_adjust)
776+
assert tip_adjust
777+
self.assertEqual(tip_adjust.as_tuple(),
778+
(-0.004453485238087255, 0.055702600244064024,
779+
-0.1626512213476071, 0.0, 0.0, 0.0))
780+
pose_with_tip_adjust = state.tip_adjust_t_base
781+
self.assertIsNotNone(pose_with_tip_adjust)
782+
assert pose_with_tip_adjust
783+
self.assertEqual(
784+
pose_with_tip_adjust.as_tuple(),
785+
(0.20666041510735556, -0.6601364983804328, 0.32705505075589714,
786+
0.11865414779070117, -3.075739605978814, -0.0166605496802987))
765787

766788

767789
class TestArm(test_utils.TestResponder):

‎pyreach/impl/host_impl.py‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -221,7 +221,8 @@ def __init__(
221221
if is_playback and block:
222222
assert isinstance(client, cli.PlaybackClient)
223223
playback_client: cli.PlaybackClient = client
224-
playback_client.next_device_data()
224+
if not playback_client.next_device_data():
225+
raise Exception("failed to read messages from the playback client")
225226
try:
226227
msg = client.get_queue().get(block=block, timeout=0.1)
227228
msgs.append(msg)
@@ -501,9 +502,21 @@ def machine_callback(
501502
self._host.start()
502503
if enable_streaming and not is_playback:
503504
self._start_streaming()
504-
else:
505+
if not is_playback:
505506
for t in self._start_arm_read_controller():
506507
t.join()
508+
else:
509+
while True:
510+
found_all_arms = True
511+
for test_arm in self._arm_devices:
512+
if test_arm.supported_controllers is None:
513+
found_all_arms = False
514+
if not found_all_arms:
515+
break
516+
assert isinstance(client, cli.PlaybackClient)
517+
playback_client_arm: cli.PlaybackClient = client
518+
if not playback_client_arm.next_device_data():
519+
raise Exception("failed to read messages from the playback client")
507520

508521
def get_timers(self) -> internal.Timers: # pylint: disable=unused-argument
509522
"""Return the global timers object."""

‎pyreach/impl/logs_directory_client.py‎

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@
1111
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
14-
1514
"""Provide a client to play back a logs directory."""
1615

1716
import json
@@ -162,6 +161,12 @@ def transform(self,
162161
if not dir_key:
163162
dir_key = dev_key
164163
msg.depth = os.path.join(self._abs_path, dir_key, filename)
164+
if not os.path.exists(msg.depth) and msg.upload_depth:
165+
abs_path, filename = os.path.split(msg.upload_depth)
166+
_, dir_key = os.path.split(abs_path)
167+
if not dir_key:
168+
dir_key = dev_key
169+
msg.depth = os.path.join(self._abs_path, dir_key, filename)
165170
return msg, utils.time_at_timestamp(msg.ts), msg.seq
166171

167172

@@ -185,15 +190,15 @@ class LogsDirectoryClient(playback_client.PlaybackClient):
185190
"""Class to implement a logs directory client."""
186191

187192
def __init__(self, robot_id: str, working_directory: str,
188-
client_id: Optional[str], select_client: bool,
193+
client_id: Optional[str], select_client_id: bool,
189194
gym_run_id: Optional[str], select_gym_run: bool) -> None:
190195
"""Init a LogsDirectoryClient.
191196
192197
Args:
193198
robot_id: the robot id to connect to.
194199
working_directory: optional directory to run within.
195200
client_id: the client ID to simulate (e.g. 00005.json vs 00000.json).
196-
select_client: if client ID is None, will select first client.
201+
select_client_id: if client ID is None, will select first client.
197202
gym_run_id: if specified, will select the given snapshot by gym_run_id.
198203
select_gym_run: if specified, but gym_run_id is None, will select first
199204
snapshot.
@@ -213,7 +218,7 @@ def __init__(self, robot_id: str, working_directory: str,
213218
device_iterator.start()
214219
command_iterator.start()
215220
self.start_playback(device_iterator, command_iterator, client_id,
216-
select_client, gym_run_id, select_gym_run, True)
221+
select_client_id, gym_run_id, select_gym_run, True)
217222
started = True
218223
finally:
219224
if not started:
@@ -222,7 +227,7 @@ def __init__(self, robot_id: str, working_directory: str,
222227

223228

224229
def connect_logs_directory(robot_id: str, working_directory: str,
225-
client_id: Optional[str], select_client: bool,
230+
client_id: Optional[str], select_client_id: bool,
226231
gym_run_id: Optional[str], select_gym_run: bool,
227232
kwargs: Dict[str, Any]) -> host.Host:
228233
"""Connect to Reach using TCP on specific host:port.
@@ -231,7 +236,7 @@ def connect_logs_directory(robot_id: str, working_directory: str,
231236
robot_id: the robot id to connect to.
232237
working_directory: optional directory to run within.
233238
client_id: the client ID to simulate.
234-
select_client: if client ID is None, will select first client.
239+
select_client_id: if client ID is None, will select first client.
235240
gym_run_id: if specified, will select the given snapshot by gym_run_id.
236241
select_gym_run: if specified, but gym_run_id is None, will select first
237242
snapshot.
@@ -241,5 +246,6 @@ def connect_logs_directory(robot_id: str, working_directory: str,
241246
Host interface if successful.
242247
"""
243248
return host_impl.HostImpl(
244-
LogsDirectoryClient(robot_id, working_directory, client_id, select_client,
245-
gym_run_id, select_gym_run), **kwargs)
249+
LogsDirectoryClient(robot_id, working_directory, client_id,
250+
select_client_id, gym_run_id, select_gym_run),
251+
**kwargs)

‎pyreach/impl/reach_tools_impl.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
},
1919
}
2020

21-
REACH_LINUX_X86_64_HASH = "036e2ff4f72db3515294ec484e394f693a56c917e4be2ee275a06abcd3be2e54"
21+
REACH_LINUX_X86_64_HASH = "c4ba5f5b4301638a9d9ec5cc40c10b04942b2806617c574c76190ac60e953f42"
2222
_REACH_SHA256 = {
2323
"linux": {
2424
"x86_64":

‎pyreach/impl/utils.py‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
# limitations under the License.
1414
"""Shared utility functions for PyReach implementation."""
1515

16+
import bz2
1617
import io
1718
import logging
1819
import time
@@ -185,7 +186,7 @@ def load_color_image_from_data(
185186
return color
186187
except PIL.UnidentifiedImageError as error:
187188
logging.warning("Unidentified (incorrect format?) image: %s", str(error))
188-
raise FileNotFoundError
189+
raise FileNotFoundError from error
189190

190191

191192
def load_depth_image_from_data(
@@ -217,6 +218,18 @@ def load_depth_image_from_data(
217218
raise FileNotFoundError
218219
depth.flags.writeable = False
219220
return depth
221+
if msg.depth.endswith(".bz2"):
222+
try:
223+
with bz2.open(msg.depth, "rb") as f:
224+
content = f.read()
225+
nparray = np.asarray(bytearray(content), dtype="uint8")
226+
depth = cv2.imdecode(nparray, cv2.IMREAD_ANYDEPTH)
227+
if depth is None:
228+
raise FileNotFoundError
229+
depth.flags.writeable = False
230+
return depth
231+
except OSError:
232+
raise FileNotFoundError from OSError
220233
depth = cv2.imread(msg.depth, cv2.IMREAD_ANYDEPTH)
221234
if depth is None:
222235
raise FileNotFoundError

‎pyreach/mock/arm_mock.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -160,8 +160,8 @@ def fetch_state(self, timeout: float = 15.0) -> Optional[arm.ArmState]:
160160
joint_angles: Tuple[float, ...] = (0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
161161
pose: core.Pose = core.Pose.from_tuple((0.0, 0.0, 0.0, 1.0, 0.0, 0.0))
162162
force: Tuple[float, ...] = (0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
163-
return arm.ArmState(time, sequence, "robot", "", joint_angles, pose, force,
164-
False, False, False, False, "", False, False,
163+
return arm.ArmState(time, sequence, "robot", "", joint_angles, pose, pose,
164+
force, False, False, False, False, "", False, False,
165165
arm.RobotMode.DEFAULT)
166166

167167
def async_fetch_state(

0 commit comments

Comments
 (0)