Create Default Assembly Rig Setup as shown in How RISE VFX Reached New Creative Heights with USD in Maya talk

maya_usd_rise_vfx_create_default_assembly_rig_setup_example.py

"""Create Default Assembly Rig Setup as shown in How RISE VFX Reached New Creative Heights with USD in Maya talk

See: https://youtu.be/8UIW-g1_heg?t=1902
"""
import pxr
import ufe
import pymel.core as pm
from mayaUsd import lib as mayaUsdLib

USD_SCHEMA = "UsdSchemaBase"


def create_default_assembly_rig_setup(assembly_rig_stage_shape, rig_grp, stage_driver_grp):
    """Create default assembly rig setup.
    
    Args:
        assembly_rig_stage_shape (pm.PyNode): Maya USD Proxy Shape to transform.
            This should reference a USD file with a default prim defined.
        rig_grp (pm.PyNode): Maya group that should act as the rig group
        stage_driver_grp (pm.PyNode): Group in which the locator should be
            placed that can be moved around.
    
    """
    stage = mayaUsdLib.GetPrim(assembly_rig_stage_shape.name(long=True)).GetStage()
    
    # create default control setup
    grp = pm.group(name='root_001_TRN', empty=True)
    pm.group(name="root_001_OFFSET")
    root_grp = pm.group(name="root_001_GRP")
    root_ctrl = pm.circle(c=(0, 0, 0), normal=(0, 1, 0), sweep=360, radius=10, degree=3, ut=0, tol=0.01, s=8, constructionHistory=False, name="root_001_CTRL")[0]
    root_ctrl.overrideEnabled.set(True)
    root_ctrl.overrideColor.set(17)
    pm.parent(root_ctrl, grp)
    
    # create default control setup
    grp = pm.group(name='rootOffset_001_TRN', empty=True)
    pm.group(name="rootOffset_001_OFFSET")
    root_offset_grp = pm.group(name="rootOffset_001_GRP")
    root_offset_ctrl = pm.circle(c=(0, 0, 0), normal=(0, 1, 0), sweep=360, radius=10, degree=3, ut=0, tol=0.01, s=8, constructionHistory=False, name="rootOffset_001_CTRL")[0]
    pm.parent(root_offset_ctrl, grp)
    pm.parent(root_offset_grp, grp)
    
    # parent control setup under rig group
    pm.parent(root_grp, rig_grp)
    pm.select(clear=True)
    
    # create the stage root driver as an locator and parent under driver group
    locator = pm.spaceLocator('root_LOC')
    pm.parent(locator, stage_driver_grp)
    pm.parentConstraint(root_offset_ctrl, locator, maintainOffset=False)
    pm.scaleConstraint(root_offset_ctrl, locator)
    # get default prim path, and build selection string to stage default prim 
    root_prim_path = str(stage.GetDefaultPrim().GetPath())
    stage_selection = "{}.{}".format(
        assembly_rig_stage_shape.name(long=True), root_prim_path
    )
    
    # create connection to default/root prim
    pm.select([stage_selection, locator])
    

def connect_maya_node_to_usd_prim(translate_op=True, rotate_op=True, scale_op=True):
    """Select two objects, one xformable prim in Maya USD Proxy and a maya native transform
    
    Then running this function will connect the maya transform to the xformable prim as
    driven transforms.
    """
    global_selection = ufe.GlobalSelection.get()
    
    ufe_object = global_selection.front()
    if USD_SCHEMA not in ufe_object.ancestorNodeTypes():
        return
    stage_path, prim_path = mayaUsdLib.proxyAccessor.getDagAndPrimFromUfe(ufe_object)
    stage_shape = pm.PyNode(stage_path)
    stage = mayaUsdLib.GetPrim(stage_path).GetStage()
    prim = stage.GetPrimAtPath(prim_path)
    
    maya_node = pm.ls(selection=True)[0]

    x_formable = pxr.UsdGeom.Xformable(prim)
    
    if translate_op and not prim.GetAttribute("xformOp:translate").IsValid():
        translate = x_formable.AddTranslateOp()
        translate.Set(pxr.Gf.Vec3f([0, 0, 0]))
    if rotate_op and not prim.GetAttribute("xformOp:rotateXYZ").IsValid():
        rotate = x_formable.AddRotateXYZOp()
        rotate.Set(pxr.Gf.Vec3f([0, 0, 0]))
    if scale_op and not prim.GetAttribute("xformOp:scale").IsValid():
        scale = x_formable.AddScaleOp()
        scale.Set(pxr.Gf.Vec3f([1, 1, 1]))

    reset_set = ["!resetXformStack!"]
    if translate_op:
        reset_set.append("xformOp:translate")
    if rotate_op:
        reset_set.append("xformOp:rotateXYZ")
    if scale_op:
        reset_set.append("xformOp:scale")
    prim.GetAttribute("xformOpOrder").Set(reset_set)
    
    for src_attr, target_attr_name in {
        maya_node.translate: "xformOp:translate",
        maya_node.rotate: "xformOp:rotateXYZ",
        maya_node.scale: "xformOp:scale",
    }.items():
        accessor = mayaUsdLib.proxyAccessor.getOrCreateAccessPlug(ufeObject=ufe_object, usdAttrName=target_attr_name)
        accessor_attr = stage_shape.attr(accessor)
        pm.connectAttr(src_attr, accessor_attr, force=True)

Export Maya USD Proxy Shape animation in layer with proxy accesor via session layer

Export the animation in a MayaUsdProxyShape layer that is done through a connected proxy accessor, as shown by Rise VFX here. This uses the fact that stepping through the timeline writes/caches the proxy accessor's time samples into the stage's session layer.

import pxr
from typing import List
from maya import cmds
import pymel.core as pm
from mayaUsd import lib as mayaUsdLib


def extract_assembly_animation(
    scene_stage_shape: pm.PyNode, 
    assembly_animation_layer_path: str,
    start_frame: int, 
    end_frame: int, 
    samples: List[float]):
):
    """Extract assembly animation from MayaUsdProxyShape.
    
    See: https://youtu.be/8UIW-g1_heg?t=2177
    
    Args:
        scene_stage_shape (pm.PyNode): The MayaUsdProxyShape node.
        assembly_animation_layer_path (str): The output path to write to.
        start_frame (int): Start frame.
        end_fram (int): End frame.
        samples (List[float]): List of float samples surrounding the frames to export,
            for examples [-0.25, 0.25] so for frame 1 we export samples 0.75, 1.0, 1.25
    
    """
    start_frame = int(start_frame)
    end_frame = int(end_frame)

    # Get stage
    scene_stage = mayaUsdLib.GetPrim(scene_stage_shape.name(long=True)).GetStage()
    export_stage = pxr.Usd.Stage.CreateNew(assembly_animation_layer_path)
    
    # Get/Create layers
    session_layer = scene_stage.GetSessionLayer()
    export_layer = export_stage.GetEditTarget().GetLayer()
    
    # Cache assembly animation
    session_layer.Clear()
    for frame in range(start_frame, end_frame + 1):
        cmds.currentTime(frame)
        for sample in samples:
            cmds.currentTime(frame + sample)
            
    # Stitch cache date to export stage layer
    pxr.UsdUtils.StitchLayers(export_layer, session_layer)
    
    session_layer.Clear()
    
    # Set stage time code
    export_stage.SetStartTimeCode(start_frame)
    export_stage.SetEndTimeCode(end_frame)
    
    export_stage.Save()

Pros and cons

Pros

• Simple hack to get the animation out, basically abusing that Maya caches the time samples to the session layer.

Cons

• Since all animated data will be cached to the session layer from the proxy accessors it means it's still non-trivial to export ONLY the data you want for potentially a subset of the stage's animation, instead of all its animated data. It depends on your use case whether you really care about that.
• Clearing the session layer will remove all session metadata, including e.g. MayaReference pulled dag paths metadata basically meaning that clearing the session layer will "break" any currently loaded references inside your USD stage. To avoid that it'd be better to only temporarily clear the layer during the export and then "revert" the original state after. This is trivial to implement however.