Fix tests and examples

Author: welinder

Makefile

@@ -27,7 +27,7 @@ test: build
 	docker run --rm -i mujoco_py pytest
 
 mount_shell:
-	docker run --rm -it -v `pwd`:/dev mujoco_py /bin/bash -c "pip uninstall -y mujoco_py; rm -rf /mujoco_py; (cd /dev; /bin/bash)"
+	docker run --rm -it -v `pwd`:/code mujoco_py /bin/bash -c "pip3 uninstall -y mujoco_py; rm -rf /mujoco_py; (cd /code; /bin/bash)"
 
 shell:
 	docker run --rm -it mujoco_py /bin/bash

examples/multigpu_rendering.py

@@ -4,78 +4,36 @@
 """
 from multiprocessing import set_start_method
 from time import perf_counter
-from mujoco_py import load_model_from_path, MjSim
-from mujoco_py.mjrenderpool import RenderProcess
-import tensorflow as tf
 
-#
-# Experiment parameters
-#
-
-# Image size for rendering
-IMAGE_WIDTH = 255
-IMAGE_HEIGHT = 255
-# Number of frames to render per sim
-N_FRAMES = 2000
-# Number of sims to run in parallel (assumes one per GPU),
-# so N_SIMS=2 assumes there are 2 GPUs available.
-N_SIMS = 2
-
-
-def setup_sim(device_id):
-    model = load_model_from_path("xmls/fetch/main.xml")
-    sim = MjSim(model)
-
-    image = sim.render(
-        IMAGE_WIDTH, IMAGE_HEIGHT, device_id=device_id)
-    assert image.shape == (IMAGE_HEIGHT, IMAGE_WIDTH, 3)
-
-    return sim
-
-
-def update_sim(sim, device_id):
-    sim.step()
-    return sim.render(IMAGE_WIDTH, IMAGE_HEIGHT, device_id=device_id)
+from mujoco_py import load_model_from_path, MjRenderPool
 
 
 def main():
-    print("main(): create processes", flush=True)
-    processes = []
-    for device_id in range(N_SIMS):
-        p = RenderProcess(
-            device_id, setup_sim, update_sim,
-            (IMAGE_HEIGHT, IMAGE_WIDTH, 3))
-        processes.append(p)
+    # Image size for rendering
+    IMAGE_WIDTH = 255
+    IMAGE_HEIGHT = 255
+    # Number of frames to render per sim
+    N_FRAMES = 2000
+    # Number of sims to run in parallel (assumes one per GPU),
+    # so N_SIMS=2 assumes there are 2 GPUs available.
+    N_SIMS = 2
 
-    for p in processes:
-        p.wait()
+    pool = MjRenderPool(load_model_from_path("xmls/tosser.xml"), device_ids=2)
 
     print("main(): start benchmarking", flush=True)
     start_t = perf_counter()
 
     for _ in range(N_FRAMES):
-        for p in processes:
-            p.update()
-
-        for p in processes:
-            p.wait()
-
-        for p in processes:
-            p.read()
+        pool.render(IMAGE_WIDTH, IMAGE_HEIGHT)
 
     t = perf_counter() - start_t
     print("Completed in %.1fs: %.3fms, %.1f FPS" % (
             t, t / (N_FRAMES * N_SIMS) * 1000, (N_FRAMES * N_SIMS) / t),
           flush=True)
 
-    print("main(): stopping processes", flush=True)
-    for p in processes:
-        p.stop()
-
     print("main(): finished", flush=True)
 
 
-print("XXX about to call main()", flush=True)
 if __name__ == "__main__":
-    set_start_method('fork')
+    set_start_method('spawn')
     main()

mujoco_py/mjrenderpool.py

@@ -122,7 +122,8 @@ def _worker_init(mjb_bytes, worker_id, device_ids,
         s.sim = MjSim(load_model_from_mjb(mjb_bytes))
 
         if modder is not None:
-            s.modder = modder(s.sim)
+            s.modder = modder(s.sim, random_state=proc_worker_id)
+            s.modder.whiten_materials()
         else:
             s.modder = None
 
@@ -142,7 +143,7 @@ def _worker_render(worker_id, state, width, height,
             s.sim.set_state(state)
             forward = True
         if randomize and s.modder is not None:
-            s.modder.rand_all()
+            s.modder.randomize()
             forward = True
         if forward:
             s.sim.forward()
@@ -162,7 +163,7 @@ def _worker_render(worker_id, state, width, height,
             device_id=s.device_id)
 
     def render(self, width, height, states=None, camera_name=None,
-               depth=False, randomize=False):
+               depth=False, randomize=False, copy=True):
         """
         Renders the simulations in batch. If no states are provided,
         the max_batch_size will be used.
@@ -175,6 +176,7 @@ def render(self, width, height, states=None, camera_name=None,
         - camera_name (str): name of camera to render from.
         - depth (bool): if True, also return depth.
         - randomize (bool): calls modder.rand_all() before rendering.
+        - copy (bool): return a copy rather than a reference
 
         Returns:
         - rgbs: NxHxWx3 numpy array of N images in batch of width W
@@ -206,11 +208,11 @@ def render(self, width, height, states=None, camera_name=None,
              for i, state in enumerate(states)])
 
         rgbs = self._shared_rgbs_array[:width * height * 3 * batch_size]
-        rgbs = rgbs.reshape(batch_size, height, width, 3)
+        rgbs = rgbs.reshape(batch_size, height, width, 3).copy()
 
         if depth:
             depths = self._shared_depths_array[:width * height * batch_size]
-            depths = depths.reshape(batch_size, height, width)
+            depths = depths.reshape(batch_size, height, width).copy()
             return rgbs, depths
         else:
             return rgbs

mujoco_py/modder.py

@@ -13,6 +13,9 @@ def __init__(self, sim, random_state=None):
         self.sim = sim
         if random_state is None:
             self.random_state = np.random.RandomState()
+        elif isinstance(random_state, int):
+            # random_state assumed to be an int
+            self.random_state = np.random.RandomState(random_state)
         else:
             self.random_state = random_state
 
@@ -308,6 +311,10 @@ def set_noise(self, name, rgb1, rgb2, fraction=0.9):
         self.upload_texture(name)
         return bitmap
 
+    def randomize(self):
+        for name in self.sim.model.geom_names:
+            self.rand_all(name)
+
     def rand_all(self, name):
         choices = [
             self.rand_checker,

mujoco_py/tests/test_render_pool.py

@@ -2,31 +2,40 @@
 import subprocess
 import sys
 
+import numpy as np
 import pytest
 
 from mujoco_py import MjSim, MjRenderPool, load_model_from_xml
-from mujoco_py.tests.utils import requires_rendering
-
+from mujoco_py.modder import TextureModder
+from mujoco_py.tests.utils import compare_imgs, requires_rendering
 
 BASIC_MODEL_XML = """
 <mujoco>
     <worldbody>
         <light name="light1" diffuse=".5 .5 .5" pos="0 0 3" dir="0 0 -1"/>
         <camera name="camera1" pos="3 0 0" zaxis="1 0 0" />
-        <geom name="geom1" pos="0.5 0.4 0.3" type="plane" size="1 1 0.1" rgba=".9 0 0 1"/>
+        <geom name="g1" pos="0.5 0.4 0.3" type="plane" size="1 1 0.1" rgba="1 1 1 1" material="m1" />
         <body pos="0 0 1" name="body1">
             <joint name="joint1" type="free"/>
-            <geom name="geom2" pos="0 1 0" type="box" size=".1 .2 .3" rgba="0 .9 0 1"/>
+            <geom name="g2" pos="0 1 0" type="box" size=".1 .2 .3" rgba="1 1 1 1" material="m2" />
             <site name="site1" pos="1 0 0" size="0.1" type="sphere"/>
             <site name="sensorsurf" pos="0 0.045 0" size=".03 .03 .03" type="ellipsoid" rgba="0.3 0.2 0.1 0.3"/>
         </body>
         <body pos="1 0 0" name="mocap1" mocap="true">
-            <geom conaffinity="0" contype="0" pos="0 0 0" size="0.01 0.01 0.01" type="box"/>
+            <geom name="g3" conaffinity="0" contype="0" pos="0 0 0" size="0.01 0.01 0.01" type="box" material="m3" rgba="1 1 1 1"/>
         </body>
     </worldbody>
     <sensor>
         <touch name="touchsensor" site="sensorsurf" />
     </sensor>
+    <asset>
+        <texture name="t1" width="33" height="36" type="2d" builtin="flat" />
+        <texture name="t2" width="34" height="39" type="2d" builtin="flat" />
+        <texture name="t3" width="31" height="37" type="2d" builtin="flat" />
+        <material name="m1" texture="t1" />
+        <material name="m2" texture="t2" />
+        <material name="m3" texture="t3" />
+    </asset>
 </mujoco>
 """
 
@@ -64,11 +73,59 @@ def mp_test_rendering():
 
     images = pool.render(100, 100)
     assert images.shape == (3, 100, 100, 3)
+    compare_imgs(images[0], 'test_render_pool.mp_test_rendering.0.png')
+    assert np.all(images[0] == images[1])
 
     images, depth = pool.render(101, 103, depth=True)
     assert images.shape == (3, 103, 101, 3)
     assert depth.shape == (3, 103, 101)
+    assert np.all(images[0] == images[1])
+    assert np.all(images[1] == images[2])
+
+
+def mp_test_cameras():
+    model = load_model_from_xml(BASIC_MODEL_XML)
+    pool = MjRenderPool(model, n_workers=1)
+
+    image = pool.render(100, 100)
+    assert image.shape == (1, 100, 100, 3)
+    compare_imgs(image[0], 'test_render_pool.mp_test_cameras.0.png')
+
+    image = pool.render(100, 100, camera_name='camera1')
+    assert image.shape == (1, 100, 100, 3)
+    compare_imgs(image[0], 'test_render_pool.mp_test_cameras.1.png')
+
+
+def mp_test_modder():
+    model = load_model_from_xml(BASIC_MODEL_XML)
+    pool = MjRenderPool(model, n_workers=2, modder=TextureModder)
+
+    images = pool.render(100, 100, randomize=True)
+    assert images.shape == (2, 100, 100, 3)
 
+    # the order of the images aren't guaranteed to be consistent
+    # between the render runs
+    images1 = pool.render(100, 100, randomize=False)
+    assert images1.shape == (2, 100, 100, 3)
+
+    if np.all(images[0] == images1[0]) and np.all(images[1] == images1[1]):
+        images_same = True
+    elif np.all(images[0] == images1[1]) and np.all(images[1] == images1[0]):
+        images_same = True
+    else:
+        images_same = False
+    assert images_same
+
+    images2 = pool.render(100, 100, randomize=True)
+    assert images2.shape == (2, 100, 100, 3)
+
+    if np.all(images[0] == images2[0]) and np.all(images[1] == images2[1]):
+        images_same = True
+    elif np.all(images[0] == images2[1]) and np.all(images[1] == images2[0]):
+        images_same = True
+    else:
+        images_same = False
+    assert not images_same
 
 def mp_test_states():
     sim = MjSim(load_model_from_xml(BASIC_MODEL_XML))
@@ -82,10 +139,15 @@ def mp_test_states():
 
     images = pool.render(100, 100, states=states[:2])
     assert images.shape == (2, 100, 100, 3)
+    compare_imgs(images[0], 'test_render_pool.mp_test_states.1.png')
+    compare_imgs(images[1], 'test_render_pool.mp_test_states.2.png')
 
     states = list(reversed(states))
     images = pool.render(100, 100, states=states)
     assert images.shape == (3, 100, 100, 3)
+    compare_imgs(images[0], 'test_render_pool.mp_test_states.3.png')
+    compare_imgs(images[1], 'test_render_pool.mp_test_states.4.png')
+    compare_imgs(images[2], 'test_render_pool.mp_test_states.5.png')
 
 
 if __name__ == '__main__':
@@ -95,3 +157,5 @@ def mp_test_states():
     mp_test_create_destroy()
     mp_test_rendering()
     mp_test_states()
+    mp_test_cameras()
+    mp_test_modder()

mujoco_py/tests/utils.py

@@ -14,7 +14,7 @@ def save_test_image(filename, array):
     Image.fromarray(array).save(filename)
 
 
-def compare_imgs(img, truth_filename):
+def compare_imgs(img, truth_filename, do_assert=True):
     """
     PROTIP: run the following to re-generate the test images:
 
@@ -30,6 +30,7 @@ def compare_imgs(img, truth_filename):
             backup_path = "%s_old%s" % (pre_path, ext)
             move(truth_filename, backup_path)
         save_test_image(truth_filename, img)
+        return 0
     true_img = np.asarray(Image.open(truth_filename))
     assert img.shape == true_img.shape
     hash0 = imagehash.dhash(Image.fromarray(img))
@@ -43,7 +44,9 @@ def compare_imgs(img, truth_filename):
         save_test_image("/tmp/img.png", img)
         save_test_image("/tmp/true_img.png", true_img)
         save_test_image("/tmp/diff_img.png", img - true_img)
-    assert diff <= 1
+    if do_assert:
+        assert diff <= 1
+    return diff
 
 
 # Skips test when RENDERING_OFF.