chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'google-deepmind:main' into usd-integration

Browse Source
This commit is contained in:
Abhishek Joshi
2023-11-15 12:38:46 -06:00
committed by GitHub
parent c37212138e 8d5966eec9
commit 1491ac15d8
133 changed files with 4003 additions and 1936 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/build_requirements.txt
+2 -2
View File
@@ -41,8 +41,8 @@ numpy==1.26.0; python_version >= '3.9' \
--hash=sha256:c78a22e95182fb2e7874712433eaa610478a3caf86f28c621708d35fa4fd6e7f \
--hash=sha256:86f737708b366c36b76e953c46ba5827d8c27b7a8c9d0f471810728e5a2fe57c \
--hash=sha256:020cdbee66ed46b671429c7265cf00d8ac91c046901c55684954c3958525dab2
pip==23.2.1 \
--hash=sha256:7ccf472345f20d35bdc9d1841ff5f313260c2c33fe417f48c30ac46cccabf5be
pip==23.3.1 \
--hash=sha256:55eb67bb6171d37447e82213be585b75fe2b12b359e993773aca4de9247a052b
PyOpenGL==3.1.7 \
--hash=sha256:a6ab19cf290df6101aaf7470843a9c46207789855746399d0af92521a0a92b7a
pytest==7.4.2 \
python/make_sdist_requirements.txt
+2 -2
View File
@@ -2,8 +2,8 @@ absl-py==2.0.0 \
--hash=sha256:9a28abb62774ae4e8edbe2dd4c49ffcd45a6a848952a5eccc6a49f3f0fc1e2f3
build==1.0.3 \
--hash=sha256:589bf99a67df7c9cf07ec0ac0e5e2ea5d4b37ac63301c4986d1acb126aa83f8f
pip==23.2.1 \
--hash=sha256:7ccf472345f20d35bdc9d1841ff5f313260c2c33fe417f48c30ac46cccabf5be
pip==23.3.1 \
--hash=sha256:55eb67bb6171d37447e82213be585b75fe2b12b359e993773aca4de9247a052b
setuptools==68.2.2 \
--hash=sha256:b454a35605876da60632df1a60f736524eb73cc47bbc9f3f1ef1b644de74fd2a
python/mujoco/CMakeLists.txt
-1
View File
@@ -265,7 +265,6 @@ target_sources(
structs_header
INTERFACE indexer_xmacro.h
indexers.h
mjdata_meta.h
structs.h
)
set_target_properties(structs_header PROPERTIES PUBLIC_HEADER structs.h)
python/mujoco/bindings_test.py
+89 -4
View File
@@ -74,10 +74,21 @@ TEST_XML_SENSOR = r"""
"""
TEST_XML_PLUGIN = r"""
<mujoco model="test">
<mujoco>
<option gravity="0 0 0"/>
<extension>
<plugin plugin="mujoco.elasticity.cable"/>
</extension>
<worldbody>
<composite type="cable" curve="s" count="41 1 1" size="1" offset="0 0 1" initial="none">
<plugin plugin="mujoco.elasticity.cable">
<config key="twist" value="1e6"/>
<config key="bend" value="1e9"/>
</plugin>
<joint kind="main" damping="2"/>
<geom type="capsule" size=".005" density="1"/>
</composite>
</worldbody>
</mujoco>
"""
@@ -1048,10 +1059,24 @@ Euler integrator, semi-implicit in velocity.
while data_instances:
d = data_instances.pop()
self.assertEqual(sys.getrefcount(d), 2)
del d
while model_instances:
m = model_instances.pop()
self.assertEqual(sys.getrefcount(m), 2)
# This test is disabled on PyPy as it uses sys.getrefcount
# However PyPy is not officially supported by MuJoCo
@absltest.skipIf(sys.implementation.name == 'pypy',
reason='requires sys.getrefcount')
def test_mjdata_holds_ref_to_model(self):
data = mujoco.MjData(mujoco.MjModel.from_xml_string('<mujoco/>'))
model = data.model
# references: one in `data.model, one in `model`, one in the temporary ref
# passed to getrefcount.
self.assertEqual(sys.getrefcount(data.model), 3)
del data
self.assertEqual(sys.getrefcount(model), 2)
def test_can_initialize_mjv_structs(self):
self.assertIsInstance(mujoco.MjvScene(), mujoco.MjvScene)
self.assertIsInstance(mujoco.MjvCamera(), mujoco.MjvCamera)
@@ -1110,6 +1135,41 @@ Euler integrator, semi-implicit in velocity.
bodyexclude=0,
geomid=geomid)
def test_mj_multi_ray(self):
nray = 3
geom1 = np.zeros(1, np.int32)
pnt = np.array([-0.3, 0, 0.1])
vec = np.array([[1, 0, 0], [0, 0, 1], [0, 0, -1]], np.float64)
dist_ex = np.array([0.2, -1, 0.1])
geom_ex = np.array([1, -1, 0])
geomid = np.zeros(nray, np.int32)
dist = np.zeros(nray, np.float64)
mujoco.mj_forward(self.model, self.data)
mujoco.mj_multiRay(
m=self.model,
d=self.data,
pnt=pnt,
vec=vec.flatten(),
geomgroup=None,
flg_static=1,
bodyexclude=-1,
geomid=geomid,
dist=dist,
nray=nray,
cutoff=mujoco.mjMAXVAL)
for i in range(0, 3):
self.assertEqual(
dist[i],
mujoco.mj_ray(
self.model, self.data, pnt, vec[i], None, 1, -1, geom1
),
)
self.assertEqual(geomid[i], geom1)
self.assertEqual(geomid[i], geom_ex[i])
self.assertAlmostEqual(dist[i], dist_ex[i])
def test_inverse_fd_none(self):
eps = 1e-6
flg_centered = 0
@@ -1287,6 +1347,34 @@ Euler integrator, semi-implicit in velocity.
self.assertEqual(data.geom(3).xpos[2], 4)
self.assertEqual(data.geom(4).xpos[2], 5)
def test_load_plugin(self):
model = mujoco.MjModel.from_xml_string(TEST_XML_PLUGIN)
data = mujoco.MjData(model)
mujoco.mj_forward(model, data)
def test_copy_mjdata_with_plugin(self):
model = mujoco.MjModel.from_xml_string(TEST_XML_PLUGIN)
data1 = mujoco.MjData(model)
self.assertIs(data1.model, model)
mujoco.mj_step(model, data1)
data2 = copy.copy(data1)
mujoco.mj_step(model, data1)
mujoco.mj_step(model, data2)
np.testing.assert_array_equal(data1.qpos, data2.qpos)
self.assertIs(data1.model, data2.model)
def test_deepcopy_mjdata_with_plugin(self):
model = mujoco.MjModel.from_xml_string(TEST_XML_PLUGIN)
data1 = mujoco.MjData(model)
self.assertIs(data1.model, model)
mujoco.mj_step(model, data1)
data2 = copy.deepcopy(data1)
mujoco.mj_step(model, data1)
mujoco.mj_step(model, data2)
np.testing.assert_array_equal(data1.qpos, data2.qpos)
self.assertIsNot(data1.model, data2.model)
self.assertNotEqual(data1.model._address, data2.model._address)
def _assert_attributes_equal(self, actual_obj, expected_obj, attr_to_compare):
for name in attr_to_compare:
actual_value = getattr(actual_obj, name)
@@ -1300,9 +1388,6 @@ Euler integrator, semi-implicit in velocity.
self.fail("Attribute '{}' differs from expected value: {}".format(
name, str(e)))
def test_load_plugin(self):
mujoco.MjModel.from_xml_string(TEST_XML_PLUGIN)
if __name__ == '__main__':
absltest.main()
python/mujoco/functions.cc
+24 -3
View File
@@ -23,6 +23,8 @@
#include <Eigen/Core>
#include <mujoco/mjxmacro.h>
#include <mujoco/mujoco.h>
#include "errors.h"
#include "function_traits.h"
#include "functions.h"
#include "private.h"
@@ -37,6 +39,7 @@ PYBIND11_MODULE(_functions, pymodule) {
namespace py = ::pybind11;
namespace traits = python_traits;
using EigenVectorI = Eigen::Vector<int, Eigen::Dynamic>;
using EigenVectorX = Eigen::Vector<mjtNum, Eigen::Dynamic>;
using EigenArrayXX = Eigen::Array<
mjtNum, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
@@ -567,7 +570,25 @@ PYBIND11_MODULE(_functions, pymodule) {
Def<traits::mj_versionString>(pymodule);
// Ray collision
Def<traits::mj_multiRay>(pymodule);
Def<traits::mj_multiRay>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d, const mjtNum(*pnt)[3],
Eigen::Ref<const EigenVectorX> vec,
std::optional<Eigen::Ref<const Eigen::Vector<mjtByte, mjNGROUP>>>
geomgroup,
mjtByte flg_static, int bodyexclude, Eigen::Ref<EigenVectorI> geomid,
Eigen::Ref<EigenVectorX> dist, int nray, mjtNum cutoff) {
if (dist.size() != nray || geomid.size() != nray) {
throw py::type_error("dist and geomid should be of size nray");
}
if (vec.size() != 3 * nray) {
throw py::type_error("vec should be of size 3*nray");
}
InterceptMjErrors(::mj_multiRay)(
m, d, &(*pnt)[0], vec.data(),
geomgroup.has_value() ? geomgroup->data() : nullptr, flg_static,
bodyexclude, geomid.data(), dist.data(), nray, cutoff);
});
Def<traits::mj_ray>(
pymodule,
[](const raw::MjModel* m, const raw::MjData* d, const mjtNum(*pnt)[3],
@@ -1367,7 +1388,7 @@ PYBIND11_MODULE(_functions, pymodule) {
}
#undef MJ_M
#define MJ_M(x) d.metadata().x
#define MJ_M(x) d.model().get()->x
#undef MJ_D
#define MJ_D(x) data->x
#define X(type, name, nr, nc) \
@@ -1379,7 +1400,7 @@ PYBIND11_MODULE(_functions, pymodule) {
}
MJDATA_ARENA_POINTERS_PRIMAL
if (d.metadata().is_dual) {
if (mj_isDual(d.model().get())) {
MJDATA_ARENA_POINTERS_DUAL
}
#undef X
python/mujoco/indexers.cc
+13 -21
View File
@@ -13,16 +13,13 @@
// limitations under the License.
#include <algorithm>
#include <memory>
#include <optional>
#include <string>
#include <string_view>
#include <utility>
#include <variant>
#include <vector>
#include "indexers.h"
#include "mjdata_meta.h"
#include "raw.h"
#include "util/crossplatform.h"
@@ -76,7 +73,6 @@ IDToName MakeIDToName(int count, IntPtr name_offsets, CharPtr names) {
// the MuJoCo category of the entity itself, e.g. nbody indicates that the
// field belongs to a body.
// T: Scalar data type of the field.
// M: Either raw::MjModel or MjDataMetadata.
//
// Args:
// base_ptr: Pointer to the first entry in the entire field.
@@ -87,35 +83,35 @@ IDToName MakeIDToName(int count, IntPtr name_offsets, CharPtr names) {
// additional dimension of the size len(qvel) of the particular joint.
// m: Used for dereferencing MjSize.
// owner: The base object whose lifetime is tied to the returned array.
template <auto MjSize, typename T, typename M>
template <auto MjSize, typename T>
py::array_t<T> MakeArray(T* base_ptr, int index, std::vector<int>&& shape,
const M& m, py::handle owner) {
const raw::MjModel& m, py::handle owner) {
int offset;
if (MjSize == &M::nq) {
if (MjSize == &raw::MjModel::nq) {
offset = m.jnt_qposadr[index];
shape.insert(
shape.begin(),
((index < m.njnt-1) ? m.jnt_qposadr[index+1] : m.nq) - offset);
} else if (MjSize == &M::nv) {
} else if (MjSize == &raw::MjModel::nv) {
offset = m.jnt_dofadr[index];
shape.insert(
shape.begin(),
((index < m.njnt-1) ? m.jnt_dofadr[index+1] : m.nv) - offset);
} else if (MjSize == &M::nhfielddata) {
} else if (MjSize == &raw::MjModel::nhfielddata) {
offset = m.hfield_adr[index];
shape.insert(shape.begin(), m.hfield_ncol[index]);
shape.insert(shape.begin(), m.hfield_nrow[index]);
} else if (MjSize == &M::ntexdata) {
} else if (MjSize == &raw::MjModel::ntexdata) {
offset = m.tex_adr[index];
shape.insert(shape.begin(), m.tex_width[index]);
shape.insert(shape.begin(), m.tex_height[index]);
} else if (MjSize == &M::nsensordata) {
} else if (MjSize == &raw::MjModel::nsensordata) {
offset = m.sensor_adr[index];
shape.insert(shape.begin(), m.sensor_dim[index]);
} else if (MjSize == &M::nnumericdata) {
} else if (MjSize == &raw::MjModel::nnumericdata) {
offset = m.numeric_adr[index];
shape.insert(shape.begin(), m.numeric_size[index]);
} else if (MjSize == &M::ntupledata) {
} else if (MjSize == &raw::MjModel::ntupledata) {
offset = m.tuple_adr[index];
shape.insert(shape.begin(), m.tuple_size[index]);
} else {
@@ -135,9 +131,7 @@ py::array_t<T> MakeArray(T* base_ptr, int index, std::vector<int>&& shape,
}
} // namespace
// M is either a raw::MjModel or MjDataMetadata.
template <typename M>
NameToIDMappings::NameToIDMappings(const M& m)
NameToIDMappings::NameToIDMappings(const raw::MjModel& m)
: body(MakeNameToID(m.nbody, m.name_bodyadr, m.names)),
jnt(MakeNameToID(m.njnt, m.name_jntadr, m.names)),
geom(MakeNameToID(m.ngeom, m.name_geomadr, m.names)),
@@ -160,9 +154,7 @@ NameToIDMappings::NameToIDMappings(const M& m)
tuple(MakeNameToID(m.ntuple, m.name_tupleadr, m.names)),
key(MakeNameToID(m.nkey, m.name_keyadr, m.names)) {}
// M is either a raw::MjModel or MjDataMetadata.
template <typename M>
IDToNameMappings::IDToNameMappings(const M& m)
IDToNameMappings::IDToNameMappings(const raw::MjModel& m)
: body(MakeIDToName(m.nbody, m.name_bodyadr, m.names)),
jnt(MakeIDToName(m.njnt, m.name_jntadr, m.names)),
geom(MakeIDToName(m.ngeom, m.name_geomadr, m.names)),
@@ -223,7 +215,7 @@ MJMODEL_VIEW_GROUPS
MJMODEL_VIEW_GROUPS
#undef XGROUP
MjDataIndexer::MjDataIndexer(raw::MjData* d, const MjDataMetadata* m,
MjDataIndexer::MjDataIndexer(raw::MjData* d, const raw::MjModel* m,
py::handle owner)
: d_(d),
m_(m),
@@ -369,7 +361,7 @@ MJMODEL_KEYFRAME
#define X(type, prefix, var, dim0, dim1) \
py::array_t<type> XGROUP::var() { \
if (!var##_.has_value()) { \
var##_.emplace(MakeArray<&MjDataMetadata::dim0>( \
var##_.emplace(MakeArray<&raw::MjModel::dim0>( \
d_->prefix##var, index_, MAKE_SHAPE(dim1), *m_, owner_)); \
} \
return *var##_; \
python/mujoco/indexers.h
+6 -11
View File
@@ -25,7 +25,6 @@
#include <absl/container/flat_hash_map.h>
#include <mujoco/mjxmacro.h>
#include "indexer_xmacro.h"
#include "mjdata_meta.h"
#include "raw.h"
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
@@ -35,9 +34,7 @@ using NameToID = absl::flat_hash_map<std::string, int>;
using IDToName = std::vector<std::string>;
struct NameToIDMappings {
// M is either a raw::MjModel or MjDataMetadata.
template <typename M>
explicit NameToIDMappings(const M& m);
explicit NameToIDMappings(const raw::MjModel& m);
NameToID body;
NameToID jnt;
@@ -63,9 +60,7 @@ struct NameToIDMappings {
};
struct IDToNameMappings {
// M is either a raw::MjModel or MjDataMetadata.
template <typename M>
explicit IDToNameMappings(const M& m);
explicit IDToNameMappings(const raw::MjModel& m);
IDToName body;
IDToName jnt;
@@ -163,7 +158,7 @@ class MjModelIndexer {
class MjDataGroupedViewsBase {
public:
MjDataGroupedViewsBase(int index, std::string_view name, raw::MjData* d,
const MjDataMetadata* m,
const raw::MjModel* m,
pybind11::handle owner)
: index_(index), name_(name), d_(d), m_(m), owner_(owner) {}
@@ -175,7 +170,7 @@ class MjDataGroupedViewsBase {
int index_;
std::string name_;
raw::MjData* d_;
const MjDataMetadata* m_;
const raw::MjModel* m_;
pybind11::handle owner_;
};
@@ -201,7 +196,7 @@ MJDATA_VIEW_GROUPS
// (e.g. a particular geom or joint) either by name or by ID.
class MjDataIndexer {
public:
MjDataIndexer(raw::MjData* d, const MjDataMetadata* m,
MjDataIndexer(raw::MjData* d, const raw::MjModel* m,
pybind11::handle owner);
#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
@@ -213,7 +208,7 @@ class MjDataIndexer {
private:
raw::MjData* d_;
const MjDataMetadata* m_;
const raw::MjModel* m_;
pybind11::handle owner_;
NameToIDMappings name_to_id_;
IDToNameMappings id_to_name_;
python/mujoco/mjdata_meta.h
-108
View File
@@ -1,108 +0,0 @@
// Copyright 2021 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.
#ifndef MUJOCO_PYTHON_MJDATA_META_H_
#define MUJOCO_PYTHON_MJDATA_META_H_
#include <cstring>
#include <memory>
#include <mujoco/mujoco.h>
#include <mujoco/mjxmacro.h>
#include "raw.h"
#include "util/crossplatform.h"
namespace mujoco::python {
namespace _impl {
template <typename T> class MjWrapper;
} // namespace _impl
#define MJDATA_METADATA \
X( int, jnt_qposadr, njnt ) \
X( int, jnt_dofadr, njnt ) \
X( int, hfield_nrow, nhfield ) \
X( int, hfield_ncol, nhfield ) \
X( int, hfield_adr, nhfield ) \
X( int, tex_height, ntex ) \
X( int, tex_width, ntex ) \
X( int, tex_adr, ntex ) \
X( int, sensor_dim, nsensor ) \
X( int, sensor_adr, nsensor ) \
X( int, numeric_adr, nnumeric ) \
X( int, numeric_size, nnumeric ) \
X( int, tuple_adr, ntuple ) \
X( int, tuple_size, ntuple ) \
X( int, name_bodyadr, nbody ) \
X( int, name_jntadr, njnt ) \
X( int, name_geomadr, ngeom ) \
X( int, name_siteadr, nsite ) \
X( int, name_camadr, ncam ) \
X( int, name_lightadr, nlight ) \
X( int, name_meshadr, nmesh ) \
X( int, name_skinadr, nskin ) \
X( int, name_hfieldadr, nhfield ) \
X( int, name_texadr, ntex ) \
X( int, name_matadr, nmat ) \
X( int, name_pairadr, npair ) \
X( int, name_excludeadr, nexclude ) \
X( int, name_eqadr, neq ) \
X( int, name_tendonadr, ntendon ) \
X( int, name_actuatoradr, nu ) \
X( int, name_sensoradr, nsensor ) \
X( int, name_numericadr, nnumeric ) \
X( int, name_textadr, ntext ) \
X( int, name_tupleadr, ntuple ) \
X( int, name_keyadr, nkey ) \
X( char, names, nnames )
// A subset of mjModel fields that are required to reconstruct an MjDataWrapper.
struct MjDataMetadata {
public:
friend class _impl::MjWrapper<raw::MjData>;
explicit MjDataMetadata(const raw::MjModel* m)
:
#define X(var) var(m->var),
MJMODEL_INTS
#undef X
#define X(dtype, var, n) \
var([](dtype* src, int len) { \
dtype* dst = new dtype[len]; \
std::memcpy(dst, src, len * sizeof(dtype)); \
return dst; \
}(m->var, m->n)),
MJDATA_METADATA
#undef X
is_dual(mj_isDual(m)) {
}
#define X(var) decltype(raw::MjModel::var) var;
MJMODEL_INTS
#undef X
#define X(type, var, n) std::shared_ptr<type[]> var;
MJDATA_METADATA
#undef X
bool is_dual;
private:
MjDataMetadata() = default;
MjDataMetadata(const MjDataMetadata& other) = default;
MjDataMetadata(MjDataMetadata&& other) = default;
};
} // namespace mujoco::python
#endif // MUJOCO_PYTHON_MJDATA_META_H_
python/mujoco/mjpython/mjpython.py
+28 -1
View File
@@ -25,6 +25,8 @@ import ctypes
import importlib.util
import os
import platform
import re
import subprocess
import sys
if platform.system() != 'Darwin':
@@ -48,7 +50,32 @@ def main(argv):
# Conda doesn't create a separate shared library for Python.
# We instead use the Python binary itself, which can be dlopened just as well.
os.environ['MJPYTHON_LIBPYTHON'] = get_executable_path()
libpython_path = get_executable_path()
os.environ['MJPYTHON_LIBPYTHON'] = libpython_path
# In some installations (e.g. CommandLineTools), the Python interpreter loads
# dylibs from @executable_path-relative paths. This will not resolve
# correctly since @executable_path will be the directory containing the
# mjpython binary when we execve. We therefore preemptively resolve all
# @executable_path-relative paths now and add them to
# DYLD_FALLBACK_LIBRARY_PATH.
libpython_dir = os.path.dirname(libpython_path)
dyld_fallback_paths = (
os.environ.get('DYLD_FALLBACK_LIBRARY_PATH', '').split(':'))
pattern = re.compile(r'@executable_path/(.+) \(offset \d+\)\Z')
otool_out = subprocess.run(
['otool', '-l', libpython_path],
capture_output=True,
check=True,
).stdout.decode()
for line in otool_out.split('\n'):
m = pattern.search(line)
if m is not None:
new_path = os.path.dirname(os.path.join(libpython_dir, m.group(1)))
if new_path not in dyld_fallback_paths:
dyld_fallback_paths.insert(0, new_path)
os.environ['DYLD_FALLBACK_LIBRARY_PATH'] = ':'.join(dyld_fallback_paths)
# argv[0] is currently the path to this script.
# Replace it with sys.executable to preserve e.g. virtualenv path.
python/mujoco/structs.cc
+88 -77
View File
@@ -37,10 +37,11 @@
#include "errors.h"
#include "function_traits.h"
#include "indexers.h"
#include "mjdata_meta.h"
#include "private.h"
#include "raw.h"
#include "serialization.h"
#include <pybind11/cast.h>
#include <pybind11/detail/common.h>
#include <pybind11/numpy.h>
#include <pybind11/operators.h>
#include <pybind11/pybind11.h>
@@ -147,7 +148,8 @@ MjVisualRgbaWrapper::MjWrapper()
X(rangefinder),
X(constraint),
X(slidercrank),
X(crankbroken) {}
X(crankbroken),
X(frustum) {}
MjVisualRgbaWrapper::MjWrapper(raw::MjVisualRgba* ptr, py::handle owner)
: WrapperBase(ptr, owner),
@@ -172,7 +174,8 @@ MjVisualRgbaWrapper::MjWrapper(raw::MjVisualRgba* ptr, py::handle owner)
X(rangefinder),
X(constraint),
X(slidercrank),
X(crankbroken) {}
X(crankbroken),
X(frustum) {}
#undef X
MjVisualRgbaWrapper::MjWrapper(const MjVisualRgbaWrapper& other)
@@ -451,7 +454,8 @@ void MjModelWrapper::Serialize(std::ostream& output) const {
WriteBytes(output, buffer.data(), model_size);
}
MjModelWrapper MjModelWrapper::Deserialize(std::istream& input) {
std::unique_ptr<MjModelWrapper> MjModelWrapper::Deserialize(
std::istream& input) {
CheckInput(input, "mjModel");
char serializationVersion = ReadChar(input);
@@ -477,7 +481,7 @@ MjModelWrapper MjModelWrapper::Deserialize(std::istream& input) {
if (!model) {
throw py::value_error("Invalid serialized mjModel.");
}
return MjModelWrapper(model);
return std::unique_ptr<MjModelWrapper>(new MjModelWrapper(model));
}
// ==================== MJCONTACT ==============================================
@@ -552,13 +556,13 @@ MjDataWrapper* MjDataWrapper::FromRawPointer(raw::MjData* m) noexcept {
}
}
MjDataWrapper::MjWrapper(const MjModelWrapper& model)
: WrapperBase(InterceptMjErrors(mj_makeData)(model.get()),
MjDataWrapper::MjWrapper(MjModelWrapper* model)
: WrapperBase(InterceptMjErrors(mj_makeData)(model->get()),
&MjDataCapsuleDestructor),
#undef MJ_M
#define MJ_M(x) model.get()->x
#define MJ_M(x) model->get()->x
#define X(dtype, var, dim0, dim1) \
var(InitPyArray(X_ARRAY_SHAPE(model.get()->dim0, dim1), ptr_->var, owner_)),
var(InitPyArray(X_ARRAY_SHAPE(model->get()->dim0, dim1), ptr_->var, owner_)),
MJDATA_POINTERS
#undef MJ_M
#define MJ_M(x) (x)
@@ -569,8 +573,9 @@ MjDataWrapper::MjWrapper(const MjModelWrapper& model)
#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)),
MJDATA_VECTOR
#undef X
metadata_(model.get()),
indexer_(ptr_, &metadata_, owner_) {
model_(model),
model_ref_(py::cast(model_)),
indexer_(ptr_, model_->get(), owner_) {
bool is_newly_inserted = false;
{
py::gil_scoped_acquire gil;
@@ -585,9 +590,9 @@ MjDataWrapper::MjWrapper(const MjModelWrapper& model)
MjDataWrapper::MjWrapper(const MjDataWrapper& other)
: WrapperBase(other.Copy(), &MjDataCapsuleDestructor),
#undef MJ_M
#define MJ_M(x) other.metadata_.x
#define MJ_M(x) other.model_->get()->x
#define X(dtype, var, dim0, dim1) \
var(InitPyArray(X_ARRAY_SHAPE(other.metadata_.dim0, dim1), ptr_->var, \
var(InitPyArray(X_ARRAY_SHAPE(other.model_->get()->dim0, dim1), ptr_->var, \
owner_)),
MJDATA_POINTERS
#undef MJ_M
@@ -599,8 +604,9 @@ MjDataWrapper::MjWrapper(const MjDataWrapper& other)
#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)),
MJDATA_VECTOR
#undef X
metadata_(other.metadata_),
indexer_(ptr_, &metadata_, owner_) {
model_(other.model_),
model_ref_(other.model_ref_),
indexer_(ptr_, model_->get(), owner_) {
bool is_newly_inserted = false;
{
py::gil_scoped_acquire gil;
@@ -615,9 +621,9 @@ MjDataWrapper::MjWrapper(const MjDataWrapper& other)
MjDataWrapper::MjWrapper(MjDataWrapper&& other)
: WrapperBase(other.ptr_, other.owner_),
#undef MJ_M
#define MJ_M(x) other.metadata_.x
#define MJ_M(x) other.model_->get()->x
#define X(dtype, var, dim0, dim1) \
var(InitPyArray(X_ARRAY_SHAPE(other.metadata_.dim0, dim1), ptr_->var, \
var(InitPyArray(X_ARRAY_SHAPE(other.model_->get()->dim0, dim1), ptr_->var, \
owner_)),
MJDATA_POINTERS
#undef MJ_M
@@ -629,8 +635,9 @@ MjDataWrapper::MjWrapper(MjDataWrapper&& other)
#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)),
MJDATA_VECTOR
#undef X
metadata_(other.metadata_),
indexer_(ptr_, &metadata_, owner_) {
model_(other.model_),
model_ref_(std::move(other.model_ref_)),
indexer_(ptr_, model_->get(), owner_) {
bool is_newly_inserted = false;
{
py::gil_scoped_acquire gil;
@@ -644,12 +651,13 @@ MjDataWrapper::MjWrapper(MjDataWrapper&& other)
other.ptr_ = nullptr;
}
MjDataWrapper::MjWrapper(MjDataMetadata&& metadata, raw::MjData* d)
: WrapperBase(d, &MjDataCapsuleDestructor),
MjDataWrapper::MjWrapper(const MjDataWrapper& other, MjModelWrapper* model)
: WrapperBase(other.Copy(), &MjDataCapsuleDestructor),
#undef MJ_M
#define MJ_M(x) metadata.x
#define X(dtype, var, dim0, dim1) \
var(InitPyArray(X_ARRAY_SHAPE(metadata.dim0, dim1), ptr_->var, owner_)),
#define MJ_M(x) other.model_->get()->x
#define X(dtype, var, dim0, dim1) \
var(InitPyArray(X_ARRAY_SHAPE(other.model_->get()->dim0, dim1), ptr_->var, \
owner_)),
MJDATA_POINTERS
#undef MJ_M
#define MJ_M(x) (x)
@@ -660,8 +668,39 @@ MjDataWrapper::MjWrapper(MjDataMetadata&& metadata, raw::MjData* d)
#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)),
MJDATA_VECTOR
#undef X
metadata_(std::move(metadata)),
indexer_(ptr_, &metadata_, owner_) {
model_(model),
model_ref_(py::cast(model_)),
indexer_(ptr_, model_->get(), owner_) {
bool is_newly_inserted = false;
{
py::gil_scoped_acquire gil;
is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second;
}
if (!is_newly_inserted) {
throw UnexpectedError(
"MjDataRawPointerMap already contains this raw mjData*");
}
}
MjDataWrapper::MjWrapper(MjModelWrapper* model, raw::MjData* d)
: WrapperBase(d, &MjDataCapsuleDestructor),
#undef MJ_M
#define MJ_M(x) model->get()->x
#define X(dtype, var, dim0, dim1) \
var(InitPyArray(X_ARRAY_SHAPE(model->get()->dim0, dim1), ptr_->var, owner_)),
MJDATA_POINTERS
#undef MJ_M
#define MJ_M(x) (x)
#undef X
contact(MjContactList(ptr_->contact, NConMax(ptr_), &ptr_->ncon, owner_)),
#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)),
MJDATA_VECTOR
#undef X
model_(model),
model_ref_(py::cast(model_)),
indexer_(ptr_, model_->get(), owner_) {
bool is_newly_inserted = false;
{
py::gil_scoped_acquire gil;
@@ -692,19 +731,7 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
// TODO: Replace this custom serialization with a protobuf
WriteChar(output, kSerializationVersion);
// Write all size fields
#define X(var) WriteInt(output, this->metadata_.var);
MJMODEL_INTS
#undef X
WriteInt(output, this->metadata_.is_dual);
#define X(dtype, var, n) \
WriteBytes(output, this->metadata_.var.get(), \
this->metadata_.n * sizeof(dtype));
MJDATA_METADATA
#undef X
model_->Serialize(output);
// Write struct and scalar fields
#define X(var) WriteBytes(output, &ptr_->var, sizeof(ptr_->var))
@@ -727,15 +754,15 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
// Write buffer contents
{
MJDATA_POINTERS_PREAMBLE((&this->metadata_))
MJDATA_POINTERS_PREAMBLE((this->model_->get()))
#define X(type, name, nr, nc) \
WriteBytes(output, ptr_->name, sizeof(type)*(this->metadata_.nr)*(nc));
WriteBytes(output, ptr_->name, sizeof(type)*(this->model_->get()->nr)*(nc));
MJDATA_POINTERS
#undef X
#undef MJ_M
#define MJ_M(x) this->metadata_.x
#define MJ_M(x) this->model_->get()->x
#undef MJ_D
#define MJ_D(x) this->ptr_->x
#define X(type, name, nr, nc) \
@@ -745,7 +772,7 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
MJDATA_ARENA_POINTERS_CONTACT
MJDATA_ARENA_POINTERS_PRIMAL
if (this->metadata_.is_dual) {
if (mj_isDual(this->model_->get())) {
MJDATA_ARENA_POINTERS_DUAL
}
if (this->ptr_->nisland) {
@@ -767,28 +794,12 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
throw py::value_error("Incompatible serialization version.");
}
// Read all size and address fields
MjDataMetadata metadata;
raw::MjModel m{0};
// Read the model that was used to create the mjData.
std::unique_ptr<MjModelWrapper> m_wrapper =
MjModelWrapper::Deserialize(input);
raw::MjModel& m = *m_wrapper->get();
#define X(var) \
metadata.var = ReadInt(input); \
CheckInput(input, "mjData"); \
m.var = metadata.var;
MJMODEL_INTS
#undef X
metadata.is_dual = ReadInt(input);
#define X(dtype, var, n) \
metadata.var.reset(new dtype[metadata.n]); \
ReadBytes(input, metadata.var.get(), metadata.n * sizeof(dtype)); \
CheckInput(input, "mjData"); \
m.var = metadata.var.get();
MJDATA_METADATA
#undef X
bool is_dual = mj_isDual(&m);
raw::MjData* d = mj_makeData(&m);
if (!d) {
@@ -839,7 +850,7 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
MJDATA_ARENA_POINTERS_CONTACT
MJDATA_ARENA_POINTERS_PRIMAL
if (metadata.is_dual) {
if (is_dual) {
MJDATA_ARENA_POINTERS_DUAL
}
if (d->nisland) {
@@ -859,18 +870,12 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
throw py::value_error("Invalid serialized mjData.");
}
return MjDataWrapper(std::move(metadata), d);
return MjDataWrapper(m_wrapper.release(), d);
}
raw::MjData* MjDataWrapper::Copy() const {
raw::MjModel m{0};
#define X(var) m.var = this->metadata_.var;
MJMODEL_INTS
#undef X
#define X(dtype, var, n) m.var = this->metadata_.var.get();
MJDATA_METADATA
#undef X
return InterceptMjErrors(mj_copyData)(NULL, &m, this->ptr_);
const raw::MjModel* m = model_->get();
return InterceptMjErrors(mj_copyData)(NULL, m, this->ptr_);
}
// ==================== MJSTATISTIC ============================================
@@ -1488,6 +1493,7 @@ PYBIND11_MODULE(_structs, m) {
X(framewidth);
X(constraint);
X(slidercrank);
X(frustum);
#undef X
py::class_<MjVisualRgbaWrapper> mjVisualRgba(mjVisual, "Rgba");
@@ -1522,6 +1528,7 @@ PYBIND11_MODULE(_structs, m) {
X(constraint);
X(slidercrank);
X(crankbroken);
X(frustum);
#undef X
#define X(var) \
@@ -1886,15 +1893,19 @@ This is useful for example when the MJB is not available as a file on disk.)"));
// ==================== MJDATA ===============================================
py::class_<MjDataWrapper> mjData(m, "MjData");
mjData.def(py::init<const MjModelWrapper&>());
mjData.def(py::init<MjModelWrapper*>());
mjData.def_property_readonly("_address", [](const MjDataWrapper& d) {
return reinterpret_cast<std::uintptr_t>(d.get());
});
mjData.def_property_readonly("model", [](const MjDataWrapper& d) {
return &d.model();
});
mjData.def("__copy__", [](const MjDataWrapper& other) {
return MjDataWrapper(other);
});
mjData.def("__deepcopy__", [](const MjDataWrapper& other, py::dict) {
return MjDataWrapper(other);
MjModelWrapper* model_copy = new MjModelWrapper(other.model());
return MjDataWrapper(other, model_copy);
});
mjData.def(py::pickle(
[](const MjDataWrapper& d) { // __getstate__
@@ -1923,7 +1934,7 @@ This is useful for example when the MJB is not available as a file on disk.)"));
#undef X
#undef MJ_M
#define MJ_M(x) d.metadata().x
#define MJ_M(x) d.model().get()->x
#undef MJ_D
#define MJ_D(x) d.get()->x
#define X(dtype, var, dim0, dim1) \
python/mujoco/structs.h
+14 -6
View File
@@ -29,10 +29,10 @@
#include <mujoco/mujoco.h>
#include <mujoco/mjxmacro.h>
#include "indexers.h"
#include "mjdata_meta.h"
#include "raw.h"
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <pybind11/pytypes.h>
namespace mujoco::python {
namespace _impl {
@@ -232,6 +232,7 @@ class MjWrapper<raw::MjVisualRgba> : public WrapperBase<raw::MjVisualRgba> {
X(constraint);
X(slidercrank);
X(crankbroken);
X(frustum);
#undef X
};
@@ -446,7 +447,8 @@ class MjWrapper<raw::MjModel> : public WrapperBase<raw::MjModel> {
MjModelIndexer& indexer() { return indexer_; }
void Serialize(std::ostream& output) const;
static MjWrapper<raw::MjModel> Deserialize(std::istream& input);
static std::unique_ptr<MjWrapper<raw::MjModel>> Deserialize(
std::istream& input);
static MjWrapper LoadXMLFile(
const std::string& filename,
@@ -563,12 +565,14 @@ struct is_mj_struct_list<raw::MjContact> {
template <>
class MjWrapper<raw::MjData>: public WrapperBase<raw::MjData> {
public:
explicit MjWrapper(const MjModelWrapper& model);
explicit MjWrapper(MjModelWrapper* model);
MjWrapper(const MjWrapper& other);
MjWrapper(MjWrapper&&);
// Used for deepcopy
MjWrapper(const MjWrapper& other, MjModelWrapper* model);
~MjWrapper();
const MjDataMetadata& metadata() const { return metadata_; }
const MjModelWrapper& model() const { return *model_; }
MjDataIndexer& indexer() { return indexer_; }
void Serialize(std::ostream& output) const;
@@ -597,10 +601,14 @@ class MjWrapper<raw::MjData>: public WrapperBase<raw::MjData> {
protected:
// Internal constructor which takes ownership of given mjData pointer.
// Used for deserialization.
explicit MjWrapper(MjDataMetadata&& metadata, raw::MjData* d);
explicit MjWrapper(MjModelWrapper* model, raw::MjData* d);
raw::MjData* Copy() const;
MjDataMetadata metadata_;
// A reference to the model that was used to create this mjData.
MjModelWrapper* model_;
// A py::object pointing to the same model as model_, to make sure Python
// doesn't doesn't garbage collect it until this mjData is released.
pybind11::object model_ref_;
MjDataIndexer indexer_;
};
python/tutorial.ipynb
+216 -4
View File
@@ -72,7 +72,7 @@
},
"outputs": [],
"source": [
"#@title Check if installation was successful\n",
"#@title Set up rendering, check installation\n",
"\n",
"from google.colab import files\n",
"\n",
@@ -921,7 +921,7 @@
"id": "g1MKUEL_eSCM"
},
"source": [
"Below is a model of a chaotic pendulum, similar to [this one](https://www.exploratorium.edu/exhibits/chaotic-pendulum) in the San Francisco Exploratorium. "
"Below is a model of a chaotic pendulum, similar to [this one](https://www.exploratorium.edu/exhibits/chaotic-pendulum) in the San Francisco Exploratorium."
]
},
{
@@ -1242,7 +1242,7 @@
"free_body_MJCF = \"\"\"\n",
"\u003cmujoco\u003e\n",
" \u003casset\u003e\n",
" \u003ctexture name=\"grid\" type=\"2d\" builtin=\"checker\" rgb1=\".1 .2 .3\" \n",
" \u003ctexture name=\"grid\" type=\"2d\" builtin=\"checker\" rgb1=\".1 .2 .3\"\n",
" rgb2=\".2 .3 .4\" width=\"300\" height=\"300\" mark=\"edge\" markrgb=\".2 .3 .4\"/\u003e\n",
" \u003cmaterial name=\"grid\" texture=\"grid\" texrepeat=\"2 2\" texuniform=\"true\"\n",
" reflectance=\".2\"/\u003e\n",
@@ -1447,7 +1447,7 @@
" \u003cworldbody\u003e\n",
" \u003clight name=\"light\" pos=\"-.2 0 1\"/\u003e\n",
" \u003cgeom name=\"ground\" type=\"plane\" size=\".5 .5 10\" material=\"grid\"\n",
" zaxis=\"-.3 0 1\" friction=\".1\"/\u003e \n",
" zaxis=\"-.3 0 1\" friction=\".1\"/\u003e\n",
" \u003ccamera name=\"y\" pos=\"-.1 -.6 .3\" xyaxes=\"1 0 0 0 1 2\"/\u003e\n",
" \u003cbody pos=\"0 0 .1\"\u003e\n",
" \u003cjoint/\u003e\n",
@@ -1941,6 +1941,218 @@
" frames.append(pixels)\n",
"media.show_video(frames, fps=framerate)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Zzzugf-qPExb"
},
"source": [
"## Camera control\n",
"\n",
"Cameras can be controlled dynamically in order to achieve cinematic effects. Run the three cells below to see the difference between rendering from a static and moving camera.\n",
"\n",
"The camera-control code smoothly transitions between two trajectories, one orbiting a fixed point, the other tracking a moving object. Parameter values in the code were obtained by iterating quickly on low-res videos."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "-SW-K9WuPGrp"
},
"outputs": [],
"source": [
"#@title Load the \"dominos\" model\n",
"dominos_xml = \"\"\"\n",
"\u003cmujoco\u003e\n",
" \u003casset\u003e\n",
" \u003ctexture type=\"skybox\" builtin=\"gradient\" rgb1=\".3 .5 .7\" rgb2=\"0 0 0\" width=\"32\" height=\"512\"/\u003e\n",
" \u003ctexture name=\"grid\" type=\"2d\" builtin=\"checker\" width=\"512\" height=\"512\" rgb1=\".1 .2 .3\" rgb2=\".2 .3 .4\"/\u003e\n",
" \u003cmaterial name=\"grid\" texture=\"grid\" texrepeat=\"2 2\" texuniform=\"true\" reflectance=\".2\"/\u003e\n",
" \u003c/asset\u003e\n",
"\n",
" \u003cstatistic meansize=\".01\"/\u003e\n",
"\n",
" \u003cvisual\u003e\n",
" \u003cglobal offheight=\"2160\" offwidth=\"3840\"/\u003e\n",
" \u003cquality offsamples=\"8\"/\u003e\n",
" \u003c/visual\u003e\n",
"\n",
" \u003cdefault\u003e\n",
" \u003cgeom type=\"box\" solref=\".005 1\"/\u003e\n",
" \u003cdefault class=\"static\"\u003e\n",
" \u003cgeom rgba=\".3 .5 .7 1\"/\u003e\n",
" \u003c/default\u003e\n",
" \u003c/default\u003e\n",
"\n",
" \u003coption timestep=\"5e-4\"/\u003e\n",
"\n",
" \u003cworldbody\u003e\n",
" \u003clight pos=\".3 -.3 .8\" mode=\"trackcom\" diffuse=\"1 1 1\" specular=\".3 .3 .3\"/\u003e\n",
" \u003clight pos=\"0 -.3 .4\" mode=\"targetbodycom\" target=\"box\" diffuse=\".8 .8 .8\" specular=\".3 .3 .3\"/\u003e\n",
" \u003cgeom name=\"floor\" type=\"plane\" size=\"3 3 .01\" pos=\"-0.025 -0.295 0\" material=\"grid\"/\u003e\n",
" \u003cgeom name=\"ramp\" pos=\".25 -.45 -.03\" size=\".04 .1 .07\" euler=\"-30 0 0\" class=\"static\"/\u003e\n",
" \u003ccamera name=\"top\" pos=\"-0.37 -0.78 0.49\" xyaxes=\"0.78 -0.63 0 0.27 0.33 0.9\"/\u003e\n",
"\n",
" \u003cbody name=\"ball\" pos=\".25 -.45 .1\"\u003e\n",
" \u003cfreejoint name=\"ball\"/\u003e\n",
" \u003cgeom name=\"ball\" type=\"sphere\" size=\".02\" rgba=\".65 .81 .55 1\"/\u003e\n",
" \u003c/body\u003e\n",
"\n",
" \u003cbody pos=\".26 -.3 .03\" euler=\"0 0 -90.0\"\u003e\n",
" \u003cfreejoint/\u003e\n",
" \u003cgeom size=\".0015 .015 .03\" rgba=\"1 .5 .5 1\"/\u003e\n",
" \u003c/body\u003e\n",
"\n",
" \u003cbody pos=\".26 -.27 .04\" euler=\"0 0 -81.0\"\u003e\n",
" \u003cfreejoint/\u003e\n",
" \u003cgeom size=\".002 .02 .04\" rgba=\"1 1 .5 1\"/\u003e\n",
" \u003c/body\u003e\n",
"\n",
" \u003cbody pos=\".24 -.21 .06\" euler=\"0 0 -63.0\"\u003e\n",
" \u003cfreejoint/\u003e\n",
" \u003cgeom size=\".003 .03 .06\" rgba=\".5 1 .5 1\"/\u003e\n",
" \u003c/body\u003e\n",
"\n",
" \u003cbody pos=\".2 -.16 .08\" euler=\"0 0 -45.0\"\u003e\n",
" \u003cfreejoint/\u003e\n",
" \u003cgeom size=\".004 .04 .08\" rgba=\".5 1 1 1\"/\u003e\n",
" \u003c/body\u003e\n",
"\n",
" \u003cbody pos=\".15 -.12 .1\" euler=\"0 0 -27.0\"\u003e\n",
" \u003cfreejoint/\u003e\n",
" \u003cgeom size=\".005 .05 .1\" rgba=\".5 .5 1 1\"/\u003e\n",
" \u003c/body\u003e\n",
"\n",
" \u003cbody pos=\".09 -.1 .12\" euler=\"0 0 -9.0\"\u003e\n",
" \u003cfreejoint/\u003e\n",
" \u003cgeom size=\".006 .06 .12\" rgba=\"1 .5 1 1\"/\u003e\n",
" \u003c/body\u003e\n",
"\n",
" \u003cbody name=\"seasaw_wrapper\" pos=\"-.23 -.1 0\" euler=\"0 0 30\"\u003e\n",
" \u003cgeom size=\".01 .01 .015\" pos=\"0 .05 .015\" class=\"static\"/\u003e\n",
" \u003cgeom size=\".01 .01 .015\" pos=\"0 -.05 .015\" class=\"static\"/\u003e\n",
" \u003cgeom type=\"cylinder\" size=\".01 .0175\" pos=\"-.09 0 .0175\" class=\"static\"/\u003e\n",
" \u003cbody name=\"seasaw\" pos=\"0 0 .03\"\u003e\n",
" \u003cjoint axis=\"0 1 0\"/\u003e\n",
" \u003cgeom type=\"cylinder\" size=\".005 .039\" zaxis=\"0 1 0\" rgba=\".84 .15 .33 1\"/\u003e\n",
" \u003cgeom size=\".1 .02 .005\" pos=\"0 0 .01\" rgba=\".84 .15 .33 1\"/\u003e\n",
" \u003c/body\u003e\n",
" \u003c/body\u003e\n",
"\n",
" \u003cbody name=\"box\" pos=\"-.3 -.14 .05501\" euler=\"0 0 -30\"\u003e\n",
" \u003cfreejoint name=\"box\"/\u003e\n",
" \u003cgeom name=\"box\" size=\".01 .01 .01\" rgba=\".0 .7 .79 1\"/\u003e\n",
" \u003c/body\u003e\n",
" \u003c/worldbody\u003e\n",
"\u003c/mujoco\u003e\n",
"\"\"\"\n",
"model = mujoco.MjModel.from_xml_string(dominos_xml)\n",
"data = mujoco.MjData(model)\n",
"renderer = mujoco.Renderer(model, height=1024, width=1440)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "a2WruafiPhPk"
},
"outputs": [],
"source": [
"#@title Render from fixed camera\n",
"duration = 2.5 # (seconds)\n",
"framerate = 60 # (Hz)\n",
"\n",
"# Simulate and display video.\n",
"frames = []\n",
"mujoco.mj_resetData(model, data) # Reset state and time.\n",
"while data.time \u003c duration:\n",
" mujoco.mj_step(model, data)\n",
" if len(frames) \u003c data.time * framerate:\n",
" renderer.update_scene(data, camera='top')\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"media.show_video(frames, fps=framerate)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "Kie3y-27bQ3J"
},
"outputs": [],
"source": [
"#@title Render from moving camera\n",
"duration = 3 # (seconds)\n",
"\n",
"# find time when box is thrown (speed \u003e 2cm/s)\n",
"throw_time = 0.0\n",
"mujoco.mj_resetData(model, data)\n",
"while data.time \u003c duration and not throw_time:\n",
" mujoco.mj_step(model, data)\n",
" box_speed = np.linalg.norm(data.joint('box').qvel[:3])\n",
" if box_speed \u003e 0.02:\n",
" throw_time = data.time\n",
"assert throw_time \u003e 0\n",
"\n",
"def mix(time, t0=0.0, width=1.0):\n",
" \"\"\"Sigmoidal mixing function.\"\"\"\n",
" t = (time - t0) / width\n",
" s = 1 / (1 + np.exp(-t))\n",
" return 1 - s, s\n",
"\n",
"def unit_cos(t):\n",
" \"\"\"Unit cosine sigmoid from (0,0) to (1,1).\"\"\"\n",
" return 0.5 - np.cos(np.pi*np.clip(t, 0, 1))/2\n",
"\n",
"def orbit_motion(t):\n",
" \"\"\"Return orbit trajectory.\"\"\"\n",
" distance = 0.9\n",
" azimuth = 140 + 100 * unit_cos(t)\n",
" elevation = -30\n",
" lookat = data.geom('floor').xpos.copy()\n",
" return distance, azimuth, elevation, lookat\n",
"\n",
"def track_motion():\n",
" \"\"\"Return box-track trajectory.\"\"\"\n",
" distance = 0.08\n",
" azimuth = 280\n",
" elevation = -10\n",
" lookat = data.geom('box').xpos.copy()\n",
" return distance, azimuth, elevation, lookat\n",
"\n",
"def cam_motion():\n",
" \"\"\"Return sigmoidally-mixed {orbit, box-track} trajectory.\"\"\"\n",
" d0, a0, e0, l0 = orbit_motion(data.time / throw_time)\n",
" d1, a1, e1, l1 = track_motion()\n",
" mix_time = 0.3\n",
" w0, w1 = mix(data.time, throw_time, mix_time)\n",
" return w0*d0+w1*d1, w0*a0+w1*a1, w0*e0+w1*e1, w0*l0+w1*l1\n",
"\n",
"# Make a camera.\n",
"cam = mujoco.MjvCamera()\n",
"mujoco.mjv_defaultCamera(cam)\n",
"\n",
"# Simulate and display video.\n",
"framerate = 60 # (Hz)\n",
"slowdown = 4 # 4x slow-down\n",
"mujoco.mj_resetData(model, data)\n",
"frames = []\n",
"while data.time \u003c duration:\n",
" mujoco.mj_step(model, data)\n",
" if len(frames) \u003c data.time * framerate * slowdown:\n",
" cam.distance, cam.azimuth, cam.elevation, cam.lookat = cam_motion()\n",
" renderer.update_scene(data, cam)\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"media.show_video(frames, fps=framerate)"
]
}
],
"metadata": {