From c79ae6f24222bd3c6c0f42dbcac9c03686baf896 Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Thu, 4 Jul 2024 15:06:25 -0500 Subject: [PATCH] fixing frame rate bug in demo file --- python/mujoco/usd/camera.py | 3 +- python/mujoco/usd/demo.py | 60 ++++++++++++------------------ python/mujoco/usd/exporter.py | 13 ++++--- python/mujoco/usd/exporter_test.py | 2 +- python/mujoco/usd/objects.py | 7 +++- python/mujoco/usd/shapes.py | 2 +- 6 files changed, 41 insertions(+), 46 deletions(-) diff --git a/python/mujoco/usd/camera.py b/python/mujoco/usd/camera.py index b1b23b14..bb9d723f 100644 --- a/python/mujoco/usd/camera.py +++ b/python/mujoco/usd/camera.py @@ -18,7 +18,8 @@ from pxr import Gf from pxr import Usd from pxr import UsdGeom -import mujoco.usd.utils as utils_component +# import mujoco.usd.utils as utils_component +import utils as utils_component class USDCamera: """Class that handles the cameras in the USD scene""" diff --git a/python/mujoco/usd/demo.py b/python/mujoco/usd/demo.py index 460ce6d9..ab5d6c7a 100644 --- a/python/mujoco/usd/demo.py +++ b/python/mujoco/usd/demo.py @@ -1,26 +1,14 @@ -# Copyright 2024 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== import argparse -from pathlib import Path from tqdm import tqdm +from pathlib import Path import mujoco -from mujoco.usd import exporter + +# from mujoco.usd import exporter +import exporter def generate_usd_trajectory(args): - """Generates a USD file from a mujoco trajectory.""" + # load a model to mujoco model_path = args.model_path m = mujoco.MjModel.from_xml_path(model_path) @@ -31,12 +19,12 @@ def generate_usd_trajectory(args): output_directory_name=Path(args.model_path).stem, output_directory_root=args.output_directory_root, camera_names=args.camera_names) - - # step through the model for length steps - for _ in tqdm(range(args.length)): - for _ in range(args.steps_per_frame): - mujoco.mj_step(m, d) - exp.update_scene(d) + + # step through the simulation for the given duration of time + while d.time < args.duration: + mujoco.mj_step(m, d) + if exp.frame_count < d.time * args.framerate: + exp.update_scene(data=d) exp.save_scene(filetype=args.export_extension) @@ -44,35 +32,35 @@ if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--model_path', + parser.add_argument('--model_path', type=str, required=True, help='path to mjcf xml model') - parser.add_argument('--length', + parser.add_argument('--duration', type=int, - default=100, - help='length of trajectory to render') + default=5, + help='duration in seconds for the generated video') - parser.add_argument('--output_directory_root', + parser.add_argument('--framerate', + type=int, + default=60, + help='frame rate of the generated video') + + parser.add_argument('--output_directory_root', type=str, default="../usd_trajectories/", help='location where to create usd files') - parser.add_argument('--camera_names', + parser.add_argument('--camera_names', type=str, nargs='+', help='cameras to include in usd') - parser.add_argument('--export_extension', + parser.add_argument('--export_extension', type=str, default="usd", - help='extension of exported file (usd, usda, or usdc)') - - parser.add_argument('--steps_per_frame', - type=int, - default=1, - help='number of frames to skip for each rendering step') + help='extension of exported file (can be usd, usda, or usdc)') args = parser.parse_args() generate_usd_trajectory(args) diff --git a/python/mujoco/usd/exporter.py b/python/mujoco/usd/exporter.py index 9eedfb60..25cd6747 100644 --- a/python/mujoco/usd/exporter.py +++ b/python/mujoco/usd/exporter.py @@ -29,10 +29,14 @@ from pxr import UsdGeom import mujoco -import mujoco.usd.shapes as shapes_module -import mujoco.usd.objects as object_module -import mujoco.usd.lights as light_module -import mujoco.usd.camera as camera_module +# import mujoco.usd.shapes as shapes_module +# import mujoco.usd.objects as object_module +# import mujoco.usd.lights as light_module +# import mujoco.usd.camera as camera_module +import shapes as shapes_module +import objects as object_module +import lights as light_module +import camera as camera_module class USDExporter: @@ -293,7 +297,6 @@ class USDExporter: self.geom_refs[geom_name] = usd_geom def _update_geoms(self): - # iterate through all geoms in the scene and makes update for i in range(self.scene.ngeom): geom = self.scene.geoms[i] diff --git a/python/mujoco/usd/exporter_test.py b/python/mujoco/usd/exporter_test.py index 17d6e855..3320c77d 100644 --- a/python/mujoco/usd/exporter_test.py +++ b/python/mujoco/usd/exporter_test.py @@ -59,7 +59,7 @@ class ExporterTest(absltest.TestCase): exporter.save_scene("export.usda") with open(os.path.join( - output_dir, "mujoco_usdpkg/frames", "frame_1_.export.usda"), "r") as f: + output_dir, "mujoco_usdpkg/frames", "frame_1.export.usda"), "r") as f: golden_path = os.path.join( epath.resource_path("mujoco"), "testdata", "usd_golden.usda") with open(golden_path, "r") as golden_file: diff --git a/python/mujoco/usd/objects.py b/python/mujoco/usd/objects.py index 854606e3..42677bd7 100644 --- a/python/mujoco/usd/objects.py +++ b/python/mujoco/usd/objects.py @@ -27,8 +27,11 @@ from pxr import Vt import mujoco -import mujoco.usd.utils as utils_component -import mujoco.usd.shapes as shapes_component +# import mujoco.usd.utils as utils_component +# import mujoco.usd.shapes as shapes_component +import utils as utils_component +import shapes as shapes_component + class USDObject(ABC): """ Abstract interface for all USD objects including meshes and primitives. diff --git a/python/mujoco/usd/shapes.py b/python/mujoco/usd/shapes.py index dec9adae..51a6905f 100644 --- a/python/mujoco/usd/shapes.py +++ b/python/mujoco/usd/shapes.py @@ -151,7 +151,7 @@ def mesh_generator( prim_mesh, mesh = None, None for shape, config in mesh_config.items(): - + if "name" in shape: continue