Sam Haves
14972d2490
PiperOrigin-RevId: 780608874 Change-Id: I333238c756fe76b371518a7cdf9452b93da3540a
132 lines
4.7 KiB
C++
132 lines
4.7 KiB
C++
// Copyright 2025 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.
|
|
|
|
#include <mujoco/experimental/usd/layer_sink.h>
|
|
|
|
#include <vector>
|
|
|
|
#include <mujoco/mujoco.h>
|
|
#include <pxr/base/gf/declare.h>
|
|
#include <pxr/base/gf/matrix4d.h>
|
|
#include <pxr/base/gf/quaternion.h>
|
|
#include <pxr/base/gf/rotation.h>
|
|
#include <pxr/base/gf/transform.h>
|
|
#include <pxr/base/tf/staticTokens.h>
|
|
#include <pxr/base/tf/token.h>
|
|
#include <pxr/base/vt/types.h>
|
|
#include <pxr/base/vt/value.h>
|
|
#include <pxr/pxr.h>
|
|
#include <pxr/usd/sdf/attributeSpec.h>
|
|
#include <pxr/usd/sdf/changeBlock.h>
|
|
#include <pxr/usd/sdf/childrenPolicies.h>
|
|
#include <pxr/usd/sdf/primSpec.h>
|
|
#include <pxr/usd/sdf/types.h>
|
|
#include <pxr/usd/sdf/valueTypeName.h>
|
|
#include <pxr/usd/usd/common.h>
|
|
#include <pxr/usd/usd/stage.h>
|
|
#include <pxr/usd/usd/timeCode.h>
|
|
#include <pxr/usd/usdGeom/tokens.h>
|
|
#include <pxr/usd/usdGeom/xformCache.h>
|
|
|
|
PXR_NAMESPACE_OPEN_SCOPE
|
|
TF_DEFINE_PRIVATE_TOKENS(_layer_sink_tokens,
|
|
((xformOpTransform, "xformOp:transform:mujoco")));
|
|
PXR_NAMESPACE_CLOSE_SCOPE
|
|
|
|
namespace {
|
|
pxr::SdfAttributeSpecHandle GetOrCreateAttribute(
|
|
pxr::SdfLayerHandle layer, pxr::SdfPrimSpecHandle prim_spec,
|
|
const pxr::SdfPath& path, const pxr::TfToken& token,
|
|
const pxr::SdfValueTypeName& type, pxr::SdfVariability variability) {
|
|
// Find or create transform attribute.
|
|
pxr::SdfAttributeSpecHandle attr =
|
|
layer->GetAttributeAtPath(path.AppendProperty(token));
|
|
|
|
if (!attr) {
|
|
attr = pxr::SdfAttributeSpec::New(prim_spec, token, type, variability);
|
|
}
|
|
return attr;
|
|
}
|
|
} // namespace
|
|
|
|
namespace mujoco {
|
|
namespace usd {
|
|
LayerSink::LayerSink(pxr::UsdStageRefPtr stage) : stage_(stage) {}
|
|
|
|
void LayerSink::Update(const mjData* const data,
|
|
std::vector<pxr::SdfPath> body_paths) {
|
|
pxr::SdfLayerHandle layer = stage_->GetEditTarget().GetLayer();
|
|
|
|
pxr::UsdTimeCode timecode = data->time * layer->GetTimeCodesPerSecond();
|
|
pxr::UsdGeomXformCache xform_cache(timecode);
|
|
|
|
pxr::SdfChangeBlock change_block;
|
|
// TODO(shaves) this can be parallelised.
|
|
for (int i = 0; i < body_paths.size(); i++) {
|
|
const pxr::SdfPath& path = body_paths[i];
|
|
|
|
// Author transform timesample
|
|
// Convert pose data.
|
|
mjtNum* xpos = &data->xpos[i * 3];
|
|
mjtNum* xquat = &data->xquat[i * 4];
|
|
pxr::GfVec3d translation(xpos[0], xpos[1], xpos[2]);
|
|
pxr::GfVec3d imag_part(xquat[1], xquat[2], xquat[3]);
|
|
// Some bodies such as the world body don't have mapped USD prims to
|
|
// transform.
|
|
if (path.IsEmpty()) {
|
|
continue;
|
|
}
|
|
|
|
pxr::SdfPrimSpecHandle prim_spec = pxr::SdfCreatePrimInLayer(layer, path);
|
|
|
|
// Find or create transform attribute.
|
|
pxr::SdfAttributeSpecHandle xform_attr = GetOrCreateAttribute(
|
|
layer, prim_spec, path, pxr::_layer_sink_tokens->xformOpTransform,
|
|
pxr::SdfValueTypeNames->Matrix4d, pxr::SdfVariabilityVarying);
|
|
|
|
pxr::GfQuaternion quat(xquat[0], imag_part);
|
|
pxr::GfRotation rotation(quat);
|
|
|
|
// Identities
|
|
pxr::GfVec3d scale(1.0);
|
|
pxr::GfVec3d pivot_position(0);
|
|
pxr::GfRotation pivot_rotation(pxr::GfQuaternion::GetIdentity());
|
|
|
|
// Set time sample
|
|
pxr::GfTransform xform(translation, rotation, scale, pivot_position,
|
|
pivot_rotation);
|
|
pxr::GfMatrix4d parent_xform =
|
|
xform_cache.GetParentToWorldTransform(stage_->GetPrimAtPath(path));
|
|
pxr::GfMatrix4d relative_xform =
|
|
xform.GetMatrix() * parent_xform.GetInverse();
|
|
layer->SetTimeSample(xform_attr->GetPath(), timecode.GetValue(),
|
|
relative_xform);
|
|
|
|
// Create xformop_order attribute if missing.
|
|
pxr::SdfAttributeSpecHandle op_order_attr = GetOrCreateAttribute(
|
|
layer, prim_spec, path, pxr::UsdGeomTokens->xformOpOrder,
|
|
pxr::SdfValueTypeNames->TokenArray, pxr::SdfVariabilityUniform);
|
|
|
|
// Override all local transforms so that this is the absolute position.
|
|
// We don't reset the xform stack here, so parent transforms will still
|
|
// affect the world pose.
|
|
const pxr::VtTokenArray xformop_order(
|
|
{pxr::_layer_sink_tokens->xformOpTransform});
|
|
|
|
op_order_attr->SetDefaultValue(pxr::VtValue(xformop_order));
|
|
}
|
|
}
|
|
} // namespace usd
|
|
} // namespace mujoco
|