merged in updates to door environments

tianheyu927 · Jul 9, 2018 · 7e05198 · 7e05198
2 parents bcb4142 + ab2c05b
commit 7e05198
Show file tree

Hide file tree

Showing 5 changed files with 267 additions and 145 deletions.
diff --git a/multiworld/core/image_env.py b/multiworld/core/image_env.py
@@ -4,7 +4,6 @@
 import warnings
 from PIL import Image
 from gym.spaces import Box, Dict
-
 from multiworld.core.wrapper_env import ProxyEnv
 
 
@@ -17,6 +16,8 @@ def __init__(
             transpose=False,
             grayscale=False,
             normalize=False,
+            reward_type='image_distance',
+            threshold=10,
     ):
         self.quick_init(locals())
         super().__init__(wrapped_env)
@@ -52,10 +53,13 @@ def __init__(
         spaces['image_desired_goal'] = img_space
         spaces['image_achieved_goal'] = img_space
         self.observation_space = Dict(spaces)
+        self.reward_type=reward_type
+        self.threshold = threshold
 
     def step(self, action):
         obs, reward, done, info = self.wrapped_env.step(action)
         new_obs = self._update_obs(obs)
+        reward = self.compute_reward(action, new_obs)
         return new_obs, reward, done, info
 
     def reset(self):
@@ -104,6 +108,12 @@ def get_goal(self):
         return goal
 
     def sample_goals(self, batch_size):
+        if self.use_goal_caching:
+            idxs = np.random.randint(0, self.num_cached_goals, batch_size)
+            goals = dict()
+            for key in self.goals.keys():
+                goals[key] = self.goals[key][idxs]
+            return goals
         if batch_size > 1:
             warnings.warn("Sampling goal images is slow")
         img_goals = np.zeros((batch_size, self.image_length))
@@ -116,8 +126,16 @@ def sample_goals(self, batch_size):
         goals['image_desired_goal'] = img_goals
         return goals
 
-    def compute_rewards(self, achieved_goals, desired_goals, info):
-        return - np.linalg.norm(achieved_goals - desired_goals, axis=1)
+    def compute_rewards(self, actions, obs):
+        achieved_goals = obs['achieved_goal']
+        desired_goals = obs['desired_goal']
+        dist = np.linalg.norm(achieved_goals - desired_goals, axis=1)
+        if self.reward_type=='image_distance':
+            return -dist
+        elif self.reward_type=='image_sparse':
+            return -(dist<self.threshold).astype(float)
+        else:
+            raise NotImplementedError()
 
 def normalize_image(image):
     assert image.dtype == np.uint8

diff --git a/multiworld/envs/assets/sawyer_xyz/door_config.xml b/multiworld/envs/assets/sawyer_xyz/door_config.xml
@@ -0,0 +1,77 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+Usage:
+
+<mujoco>
+	<compiler meshdir="../meshes/sawyer" ...></compiler>
+	<include file="shared_config.xml"></include>
+    (new stuff)
+	<worldbody>
+		<include file="sawyer_xyz_base.xml"></include>
+        (new stuff)
+	</worldbody>
+</mujoco>
+-->
+
+<mujoco>
+    <asset>
+        <mesh name="pedestal" file="pedestal.stl" />
+        <mesh name="base" file="base.stl" />
+        <mesh name="l0" file="l0.stl" />
+        <mesh name="head" file="head.stl" />
+        <mesh name="l1" file="l1.stl" />
+        <mesh name="l2" file="l2.stl" />
+        <mesh name="l3" file="l3.stl" />
+        <mesh name="l4" file="l4.stl" />
+        <mesh name="l5" file="l5.stl" />
+        <mesh name="l6" file="l6.stl" />
+    </asset>
+
+    <visual>
+        <headlight ambient="0.5 0.5 0.5"  />
+    </visual>
+
+    <compiler inertiafromgeom="auto" angle="radian" coordinate="local" eulerseq="XYZ" meshdir="../meshes/sawyer"/>
+    <size njmax="6000" nconmax="6000"/>
+    <option
+            gravity="0 0 -9.81"
+            iterations="50"
+            integrator="Euler"
+            timestep="0.003"
+            cone="pyramidal"
+            impratio="1"
+            tolerance="1e-10"
+    />
+
+    <default>
+        <joint limited="true"
+               damping="1"
+               stiffness="0"
+               armature=".1"
+               user="0"/>
+        <geom solref="0.008 1"
+              solimp="0.95 0.95 0.01"
+              user="0"
+              />
+        <motor ctrllimited="true" ctrlrange="-1 1"/>
+        <position ctrllimited="true"/>
+
+        <default class='0'>
+            <geom friction="1 0.005 0.001" rgba=".75 .6 .5 1"
+                  contype="1" conaffinity="0"/>
+        </default>
+
+        <default class='1'>
+            <geom friction="1 0.005 0.001" rgba=".7 .6 .55 1"
+                  contype="4" conaffinity="2"/>
+        </default>
+
+        <default class='arm'>
+            <joint damping="5"/>
+        </default>
+    </default>
+
+    <equality>
+        <weld body1="mocap" body2="leftclaw" solref="0.005 1"></weld>
+    </equality>
+</mujoco>
diff --git a/multiworld/envs/assets/sawyer_xyz/sawyer_door.xml b/multiworld/envs/assets/sawyer_xyz/sawyer_door.xml
@@ -1,51 +1,6 @@
 <?xml version="1.0" encoding="utf-8"?>
 <mujoco model="sawyer">
-    <compiler angle="radian" meshdir="../meshes/sawyer"/>
-    <!-- <option gravity="0 0 0" /> -->
-    <!-- <option impratio="0.01" /> -->
-    <!-- <option collision="predefined" /> -->
-    <option integrator="Euler" timestep="0.002" iterations="50"/>
-    <size njmax="500" nconmax="100" />
-    <asset>
-        <mesh name="pedestal" file="pedestal.stl" />
-        <mesh name="base" file="base.stl" />
-        <mesh name="l0" file="l0.stl" />
-        <mesh name="head" file="head.stl" />
-        <mesh name="l1" file="l1.stl" />
-        <mesh name="l2" file="l2.stl" />
-        <mesh name="l3" file="l3.stl" />
-        <mesh name="l4" file="l4.stl" />
-        <mesh name="l5" file="l5.stl" />
-        <mesh name="l6" file="l6.stl" />
-    </asset>
-
-    <visual>
-        <headlight ambient="0.5 0.5 0.5"  />
-    </visual>
-
-    <default>
-        <joint limited="true" damping="1" stiffness="0" armature=".1" user="0"/>
-        <geom solref="0.008 1" solimp="0.95 0.95 0.01" margin="0.001" user="0"/>
-        <motor ctrllimited="true" ctrlrange="-1 1"/>
-        <position ctrllimited="true"/>
-
-        <default class='0'>
-            <geom friction="1 0.005 0.001" rgba=".75 .6 .5 1" contype="1" conaffinity="0"/>
-        </default>
-
-        <default class='1'>
-            <geom friction="1 0.005 0.001" rgba=".7 .6 .55 1" contype="4" conaffinity="2"/>
-        </default>
-
-        <default class='arm'>
-            <joint damping="5"/>
-        </default>
-    </default>
-
-    <equality>
-        <weld body1="mocap" body2="leftclaw" solref="0.01"></weld>
-    </equality>
-
+   <include file="door_config.xml"></include>
     <worldbody>
         <camera pos="0 0 3.2" name="topview" />
         <!-- <light directional="true" cutoff="40" exponent="10" diffuse=".4 .4 .4" specular=".7 .7 .7" pos="0 0 1.5" dir="0 0 -1"/> -->
@@ -57,9 +12,10 @@
 
         <geom name="tableTop3" type="plane" pos="-0.6 0 0" size="0.4 0.4 0.5" rgba=".6 .6 .5 1" contype="1" conaffinity="1" />
 
-	    <body name="cabinet" pos="-.775 -.05 0.15">
+        <body name="cabinet" pos="-.775 -.05 0.15">
 	        <body name="door" pos="0.8 .7. -.5">
 	            <joint name="doorjoint" pos="-.1 0 0" axis="0 0 1" type="hinge" limited="true" range="-1.5708 1.5708" damping="0.1" armature="0" />
+	            <!--<geom type="box" pos="0 0 0" size="0.01 0.2 0.2" rgba=".1 .1 .9 1" contype="1" conaffinity="1" quat="0.5 0.5 -0.5 -0.5"/>-->
 	            <geom type="box" pos="0 0 0" size="0.01 0.2 0.2" rgba=".1 .1 .9 1" contype="1" conaffinity="1" quat="0.5 0.5 -0.5 -0.5"/>
 	           	<geom type="box" pos="0.1 -0.15 0" size="0.05 0.01 0.01" rgba="0 0 0 1" contype="1" conaffinity="1" />
 	           	<geom type="box" pos="0.15 -0.075 0" size="0.01 0.075 0.01" rgba="0 0 0 1" contype="1" conaffinity="1"/>

diff --git a/multiworld/envs/mujoco/sawyer_reach_torque/sawyer_reach_torque_env.py b/multiworld/envs/mujoco/sawyer_reach_torque/sawyer_reach_torque_env.py
@@ -191,10 +191,6 @@ def init_angles(self):
             1.00000000e+00, 3.76632959e-14, 1.36837913e-11, 1.56567415e-23
         ]
 
-    @property
-    def goal_dim(self):
-        return 3
-
     @property
     def endeff_id(self):
         return self.model.body_names.index('leftclaw')